import random
import math

from tinygrad.nn import Linear
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import Adam

HIDDEN_LAYERS = 10
HIDDEN_WIDTH = 10 # Measure accuracy as a function of this.
BIAS = True

class TwoDimensionalClassifier:
    def __init__(self):
        self.layers = [Linear(HIDDEN_WIDTH, HIDDEN_WIDTH, bias=BIAS) for _ in range(0, HIDDEN_LAYERS)]
        self.layers = [Linear(2, HIDDEN_WIDTH, bias=BIAS)] + self.layers + [Linear(HIDDEN_WIDTH, 2, bias=BIAS)]

    def transform(self, x):
        for layer in self.layers[:-1]:
            x = layer(x)
            x = x.relu()
        return self.layers[-1](x)

    def transformLayerByLayer(self, x):
        for layer in self.layers[:-1]:
            x = layer(x).relu()
            yield x
        yield self.layers[-1](x).softmax()

    def __call__(self, x):
        return self.transform(x).softmax()

def trained(distribution, training_steps=3000):
    Tensor.training = True
    net = TwoDimensionalClassifier()
    opt = Adam([layer.weight for layer in net.layers])
    for _ in range(training_steps):
        point, classification = distribution.sample()
        guess = net(point)
        error = guess.sparse_categorical_crossentropy(classification)
        opt.zero_grad()
        error.backward()
        opt.step()

    Tensor.training = False
    return net

# A circle at the origin. Points inside the circle are class 1, points
# outside the circle are class 0.
class CircleAtOrigin:
    def __init__(self, radius=1):
        self.radius = 1

    # Gets a random point in a 3x3 square centered on the origin and its
    # classification ([0] or [1] for inside the circle).
    def sample(self):
        x = random.random() * self.radius*3 - self.radius*1.5
        y = random.random() * self.radius*3 - self.radius*1.5
        point = Tensor([x, y])
        classification = Tensor([1]) if x*x + y*y <= self.radius*self.radius else Tensor([0])
        return point, classification