import config, mystring, vector, util
import algo2/wed, data

static int def algo2_per_feature_vars(Algo2Set *set):
    """
    For a very simple classifier, minimum Euclidean distance weighted by the
    feature variance is used. This helper function calculates the variance
    for each such feature from all the training samples.
    """

    # Detect number of features.
    Algo2Sample *first_sample = data_get_first_sample(set)
    int fcount = first_sample->bitcode->h

    # Create mean and variance vector for each class.
    for int jc = 0; jc < set->classes->count; jc++:
        Algo2Class *c = *((Algo2Class **)vector_get(set->classes, jc))
         c->per_feature_means = ALLOCATE0(fcount * sizeof(V))
         c->per_feature_vars = ALLOCATE0(fcount * sizeof(V))

    set->per_feature_means = ALLOCATE0(fcount * sizeof(V))
    for int i = 0; i < fcount; i++:
        V mean = 0
        int n = 0
        for int jc = 0; jc < set->classes->count; jc++:
            Algo2Class *c = *((Algo2Class **)vector_get(set->classes, jc))
            V cmean = 0
            int cn = c->training_samples->count
            for int js = 0; js < cn; js++:
                Algo2Sample *s = *((Algo2Sample **)vector_get(
                    c->training_samples, js))
                V x = *((V *)s->bitcode->bits + i)
                cmean += x
            c->per_feature_means[i] = cn ? cmean / cn : 0
            mean += cmean
            n += cn
        set->per_feature_means[i] = n ? mean / n : 0

    set->per_feature_vars = ALLOCATE0(fcount * sizeof(V))
    for int i = 0; i < fcount; i++:
        V variance = 0
        int n = 0
        V mean = set->per_feature_means[i]
        for int jc = 0; jc < set->classes->count; jc++:
            Algo2Class *c = *((Algo2Class **)vector_get(set->classes, jc))
            V cvariance = 0
            V cmean = c->per_feature_means[i]
            int cn = c->training_samples->count
            for int js = 0; js < cn; js++:
                Algo2Sample *s = *((Algo2Sample **)vector_get(
                    c->training_samples, js))
                V x = *((V *)s->bitcode->bits + i)
                variance += (x - mean) * (x - mean)
                cvariance += (x - cmean) * (x - cmean)
                n++
            c->per_feature_vars[i] = cn ? cvariance / cn : 0
        set->per_feature_vars[i] = n ? variance / n : 0

    return fcount

V def algo2_get_score(Algo2Sample *s, Algo2Class *c):
    """
    Given a sample, get the score for a class. Score is a distance, the smaller,
    the better the sample fits into the class.
    """
    V distance = 0
    bool all0 = True

    # Determine a weight factor for each feature
    V weights[s->bitcode->h]
    for int j = 0; j < s->bitcode->h; j++:
        V v = c->per_feature_vars[j]
        if v: all0 = False
        weights[j] = 1.0 / sqrt(v)

    Vector *v1 = vector_new(sizeof(V))
    Vector *v2 = vector_new(sizeof(V))

    # Now calculate WED of all features averaged over all training samples.
    # Note: Using test samples instead of training samples here will lead to
    # a nice recognition rate of 100%, as the sample itself will be in
    # there :)
    for int j = 0; j < c->training_samples->count; j++:
        Algo2Sample *ts = *((Algo2Sample **)vector_get(c->training_samples, j))
        v1->data = ts->bitcode->bits
        v1->count = ts->bitcode->h
        v2->data = s->bitcode->bits
        v2->count = s->bitcode->h
        distance += WED(v1, v2, all0 ? None : weights)
    distance /= c->training_samples->count

    v1->data = None
    v2->data = None
    vector_del(v1)
    vector_del(v2)

    return distance

Algo2Class *def algo2_identify(Algo2Set *set, Algo2Sample *s):
    """
    Given the training set, identify a sample from the test set.
    """
    V min_distance = 0
    Algo2Class *winner = None
    for int i = 0; i < set->classes->count; i++:
        Algo2Class *c = *((Algo2Class **)vector_get(set->classes, i))

        V distance = algo2_get_score(s, c)

        if winner == None or distance < min_distance:
            min_distance = distance
            winner = c
    # printf("Winner for %s: %s (%f)\n", s->name, winner->name, min_distance)
    return winner

def algo2_test(Algo2Set *set):
    """
    Try to identify each sample in the test set.
    """
    int correct = 0
    int wrong = 0
    for int i = 0; i < set->classes->count; i++:
        Algo2Class *c = *((Algo2Class **)vector_get(set->classes, i))
        for int j = 0; j < c->test_samples->count; j++:
            Algo2Sample *ts = *((Algo2Sample **)vector_get(c->test_samples, j))
            Algo2Class *result = algo2_identify(set, ts)
            if result == c:
                correct++
            else:
                wrong++
    printf("%d correctly classified, %d not (%.1f %%).\n", correct, wrong,
        100.0 * correct / (correct + wrong))

def algo2_train(Algo2Set *set):
    algo2_per_feature_vars(set)
    #for int i = 0;  i < n; i++:
    #    printf("%d: %f %f\n", i,
    #        set->per_feature_means[i],
    #        sqrt(set->per_feature_vars[i]))

def algo2_verify(Algo2Set *set):
    """
    Try to verify each sample in the test set against all classes. We cumulate
    the scoring results for inter-class and intra-class matches to get
    FMR and FNMR curves.
    """
    int total = data_test_samples_count(set)
    int enrolled = data_training_samples_count(set)
    printf("Comparing %d images to %d enrolled.\n", total, enrolled)

    set->bins_count = 512
    for int k = 0; k < set->bins_count; k++:
        set->inter_bins[k] = 0
        set->intra_bins[k] = 0

    for int i = 0; i < set->classes->count; i++:
        Algo2Class *c = *((Algo2Class **)vector_get(set->classes, i))
        for int j = 0; j < c->test_samples->count; j++:
            Algo2Sample *ts = *((Algo2Sample **)vector_get(c->test_samples, j))
            for int i2 = 0; i2 < set->classes->count; i2++:
                Algo2Class *c2 = *((Algo2Class **)vector_get(set->classes, i2))
                # c: class of our test sample
                # ts: our test sample
                # c2: class we verify against
                V distance = algo2_get_score(ts, c2)
                int k = distance * set->bins_count
                if k > set->bins_count - 1: k = set->bins_count - 1
                if c == c2:
                    set->intra_bins[k]++
                else:
                    set->inter_bins[k]++
