from typing import Tuple, List, Callable, Generator
from import_data import get_test_data_generator, get_training_data_generator, IMAGE_SIZE, print_img_to_console
import numpy as np


class MulticlassPerceptron:
    class Perceptron:
        def __init__(self, input_size, learn_rate: float = 0.0001):
            self.weights = np.random.rand(input_size + 1) * 2 - 1
            self.learn_rate = learn_rate

        def train(self, data_lbl_pairs: Callable[[], Generator], positive_label: int, iterations: int = 10):
            for num_iter in range(iterations):
                print(f"Iteration number: {num_iter + 1}")
                for i, (x, y) in enumerate(data_lbl_pairs()):
                    prediction = self.output(x)
                    label = 1 if y == positive_label else -1
                    correct = prediction >= 0 and label >= 0 or prediction < 0 and label < 0
                    if not correct:
                        self.weights = self.weights + self.learn_rate * label * np.array([1.0] + x)

        def output(self, datum):
            return np.dot(self.weights, np.array([1.0] + datum))

        def get_normalised_weight_array(self):
            return [int(x) for x in (((self.weights / np.max(self.weights) + 1) / 2) * 255)]

    def __init__(self, input_size: int, num_classes: int, learn_rate: float = 0.001):
        self.classifiers = [self.Perceptron(input_size, learn_rate) for _ in range(num_classes)]

    def train(self, data_lbl_pairs: Callable[[], Generator], iterations: int = 10):
        for i, classifier in enumerate(self.classifiers):
            print(f"Training classifier for class {i}...")
            classifier.train(data_lbl_pairs, i, iterations)
            print()

    def output(self, datum: List[int]):
        return list(map(lambda x: x.output(datum), self.classifiers))

    def prediction(self, datum):
        return max(list(range(10)), key=lambda x: self.output(datum)[x])

    def view_for_classifier(self, classifier_index: int):
        return self.classifiers[classifier_index].get_normalised_weight_array()


def train_and_test_multiclass_perceptron(iterations: int = 5, training_inputs: int = -1, test_inputs: int = -1):
    print("Loading data")
    training_data_gen = get_training_data_generator(training_inputs)
    print("Begin training model!")
    model = MulticlassPerceptron(IMAGE_SIZE, 10)
    model.train(training_data_gen, iterations)
    print("Model successfully trained.")
    print("Testing model...")
    test_data = list(get_test_data_generator(test_inputs)())
    n_correct = sum(model.prediction(x) == y for x, y in test_data)
    accuracy = n_correct / len(test_data)
    print(f"Accuracy: {accuracy} ({n_correct} correctly classified out of {len(test_data)} total test inputs.)")
    for i in range(10):
        print_img_to_console(model.view_for_classifier(i))


def make_trained_digit_model(iterations: int = 5, training_inputs: int = 5000, test_inputs: int = 1000):
    training_data_gen = get_training_data_generator(training_inputs)
    model = MulticlassPerceptron(IMAGE_SIZE, 10)
    model.train(training_data_gen, iterations)
    return model


if __name__ == "__main__":
    train_and_test_multiclass_perceptron()