import numpy as np
import time
import os

import matplotlib.pyplot as plt
import uctc.nn as nn 
from utils import parameter_data, Dataset

use_graphics = False
class PerceptronModel(object):
    def __init__(self, dimensions):
        """
        Initialize a new Perceptron instance.

        A perceptron classifies data points as either belonging to a particular
        class (+1) or not (-1). `dimensions` is the dimensionality of the data.
        For example, dimensions=2 would mean that the perceptron must classify
        2D points.
        """
        self.w = nn.Parameter(parameter_data(dimensions, 1))

    def get_weights(self):
        """
        Return a Parameter instance with the current weights of the perceptron.
        """
        return self.w.data()

    def run(self, x):
        """
        Calculates the score assigned by the perceptron to a data point x.

        Inputs:
            x: a node with shape (1 x dimensions)
        Returns: a node containing a single number (the score)
        """
        "*** YOUR CODE HERE ***"
        out = nn.Linear(x, self.w)
        return out

    def get_prediction(self, x):
        """
        Calculates the predicted class for a single data point `x`.

        Returns: 1 or -1
        """
        "*** YOUR CODE HERE ***"
        score = self.run(x).data()[0]
        # score = np.array(x.data()).dot(np.array(self.w.data()))
        if score >= 0:
            return 1
        else:
            return -1


    def train(self, dataset):
        """
        Train the perceptron until convergence.
        """
        "*** YOUR CODE HERE ***"
        batch_size = 1

        while True:
            converged = True
            for x, y in dataset.iterate_once(batch_size):
                prediction = self.get_prediction(x)
                x = np.array(x.data(), dtype=np.float32)
                y = int(y.data()[0])
                # assert 0
                if prediction != y:
                    # print(prediction, y)
                    converged = False
                    self.w.update(nn.pyarray_to_tensor(x), -y)
                # time.sleep(0.01)
            if converged:
                break

class PerceptronDataset(Dataset):
    def __init__(self, model: PerceptronModel):
        points = 500
        x = np.hstack([np.random.randn(points, 2), np.ones((points, 1))])
        y = np.where(x[:, 0] + 2 * x[:, 1] - 1 >= 0, 1.0, -1.0)
        super().__init__(x, np.expand_dims(y, axis=1))

        self.model = model
        self.epoch = 0
        limits = np.array([-3.0, 3.0])
        if use_graphics:
            fig, ax = plt.subplots(1, 1)
            ax.set_xlim(limits)
            ax.set_ylim(limits)
            positive = ax.scatter(*x[y == 1, :-1].T, color="red", marker="+")
            negative = ax.scatter(*x[y == -1, :-1].T, color="blue", marker="_")
            line, = ax.plot([], [], color="black")
            text = ax.text(0.03, 0.97, "", transform=ax.transAxes, va="top")
            ax.legend([positive, negative], [1, -1])
            plt.show(block=False)

            self.fig = fig
            self.line = line
            self.text = text
        self.limits = limits
        self.last_update = time.time()

    def iterate_once(self, batch_size):
        self.epoch += 1

        for i, (x, y) in enumerate(super().iterate_once(batch_size)):
            yield x, y

            if time.time() - self.last_update > 0.001:
                w = self.model.get_weights()
                limits = self.limits
                print(f"epoch: {self.epoch}\npoint: {i * batch_size + 1}/{len(self.x)}\nweights: {w}")
                if use_graphics:
                    if w[1] != 0:
                        self.line.set_data(limits, (-w[0] * limits - w[2]) / w[1])
                    elif w[0] != 0:
                        self.line.set_data(np.full(2, -w[2] / w[0]), limits)
                    else:
                        self.line.set_data([], [])
                        self.text.set_text(
                                f"epoch: {self.epoch}\npoint: {i * batch_size + 1}/{len(self.x)}\nweights: {w}")
                        self.fig.canvas.draw_idle()
                        self.fig.canvas.start_event_loop(1e-3)
                self.last_update = time.time()

model = PerceptronModel(3)
dataset = PerceptronDataset(model)
model.train(dataset)