import numpy as np

import practice_check
from practice_check import check, clear_progress, equal, fixture, fixtures, progress, shape, true
from support import make_spirals

BUDGET = 1500
TARGET = 0.95
HELD_OUT_SEED = 52903
ARM_T = np.linspace(1.5, 12.0, 24)
ARM_RADIUS = 3.0 * ARM_T / (4.0 * np.pi)
ARMS = np.concatenate([np.stack([ARM_RADIUS * np.cos(ARM_T + np.pi * k), ARM_RADIUS * np.sin(ARM_T + np.pi * k)], axis=1) for k in (0, 1)])
ARM_LABELS = np.repeat([0, 1], 24)
_runs = {}


def train_data():
    return make_spirals(400, np.random.default_rng(0))


def held_out():
    return make_spirals(2000, np.random.default_rng(HELD_OUT_SEED))


def fresh_run(sub, X, y, seed):
    saved = list(progress())
    clear_progress()
    try:
        model = sub.train(X, y, np.random.default_rng(seed))
        points = list(progress())
    finally:
        clear_progress()
        practice_check._progress.extend(saved)
    return model, points


def main_run(sub):
    key = id(sub)
    if key not in _runs:
        X, y = train_data()
        X_in, y_in = X.copy(), y.copy()
        clear_progress()
        model = sub.train(X_in, y_in, np.random.default_rng(1))
        _runs[key] = {"model": model, "points": list(progress()), "X": X, "y": y, "X_in": X_in, "y_in": y_in}
    return _runs[key]


def accuracy(sub, model, X, y):
    labels = np.asarray(sub.predict(model, X))
    shape(labels, y.shape, what="the labels from predict")
    return float(np.mean(labels == y))


@check("examples", "predict gives one label for each point", "predict returns an int array of shape (N,) with the value 0 or 1 for each row of X.")
def example_shape(sub):
    run = main_run(sub)
    labels = np.asarray(sub.predict(run["model"], run["X"]))
    shape(labels, (400,), what="the labels for the 400 training points")
    true(set(np.unique(labels).tolist()) <= {0, 1}, "predict returned a value that is not 0 or 1", "only 0 and 1", np.unique(labels))


@check("examples", "the model fits the training points", "After training, at least 95 percent of the training points must get the correct label.")
def example_train_accuracy(sub):
    run = main_run(sub)
    acc = accuracy(sub, run["model"], run["X"], run["y"])
    true(acc >= TARGET, f"training accuracy {acc:.3f} is below {TARGET}", f">= {TARGET}", acc)


@check("properties", "training stays inside the step budget", "Call report(step, loss) once for each update step, with step = 0, 1, 2 and so on. Use at most 1500 steps.")
def budget(sub):
    points = main_run(sub)["points"]
    true(len(points) >= 1, "train did not call report(step, loss)", "at least 1 call", 0)
    true(len(points) <= BUDGET, f"train used {len(points)} steps, the budget is {BUDGET}", f"<= {BUDGET}", len(points))
    steps = [point[0] for point in points]
    equal(steps, list(range(len(points))), what="the step numbers given to report")


@check("properties", "the reported loss goes down", "Gradient descent must lower the loss. The mean of the last 10 reported losses must be below the first loss.")
def loss_goes_down(sub):
    losses = np.array([point[1] for point in main_run(sub)["points"]])
    true(losses.size >= 1 and np.all(np.isfinite(losses)), "a reported loss is NaN or infinite", "finite losses", losses)
    true(np.mean(losses[-10:]) < losses[0], "the loss did not go down", f"below {losses[0]:.4f}", float(np.mean(losses[-10:])))


@check("properties", "the same seed gives the same model", "Take every random number from the rng argument. Then two runs with the same seed give the same predictions.")
def deterministic(sub):
    run = main_run(sub)
    X, y = train_data()
    model, points = fresh_run(sub, X, y, 1)
    X_test, _ = held_out()
    equal(np.asarray(sub.predict(model, X_test)), np.asarray(sub.predict(run["model"], X_test)), what="the predictions of the second run")
    equal([point[1] for point in points], [point[1] for point in run["points"]], what="the losses of the second run")


@check("properties", "the inputs are not changed", "Do not write into X or y. Make a new array when you change the data.")
def inputs_unchanged(sub):
    run = main_run(sub)
    equal(run["X_in"], run["X"], what="X after train")
    equal(run["y_in"], run["y"], what="y after train")


@check("edge", "points on the center line of each arm", "Points with no noise on the center line of an arm must get the label of that arm. At least 44 of the 48 points must be correct, also on the inner turns where the arms are close.")
def arm_center_lines(sub):
    labels = np.asarray(sub.predict(main_run(sub)["model"], ARMS))
    shape(labels, ARM_LABELS.shape, what="the labels of 48 points on the two arms")
    correct = int(np.sum(labels == ARM_LABELS))
    true(correct >= 44, f"only {correct} of 48 points on the arms get the correct label", ">= 44", correct)


@check("edge", "predict on one point", "predict must work for a single row and return shape (1,).")
def one_point(sub):
    labels = np.asarray(sub.predict(main_run(sub)["model"], ARMS[5:6]))
    shape(labels, (1,), what="the label of one point")


@check("edge", "a smaller data set and a new seed", "Training must not depend on one lucky seed. With 300 points and another seed, the model must also reach the target.")
def other_seed(sub):
    X, y = make_spirals(300, np.random.default_rng(5))
    model, _ = fresh_run(sub, X, y, 7)
    X_test, y_test = held_out()
    acc = accuracy(sub, model, X_test, y_test)
    true(acc >= TARGET, f"held-out accuracy {acc:.3f} is below {TARGET}", f">= {TARGET}", acc)


@fixture("held_out")
def reference_score(ref):
    X, y = train_data()
    model = ref.train(X, y, np.random.default_rng(1))
    X_test, y_test = held_out()
    yield {"target": TARGET}, float(np.mean(ref.predict(model, X_test) == y_test))


@check("reference", "held-out accuracy reaches the target", "The checks measure your model on 2000 new points that train never sees. The accuracy must be at least 0.95.")
def held_out_accuracy(sub):
    (case, reference), = fixtures("held_out")
    run = main_run(sub)
    X_test, y_test = held_out()
    acc = accuracy(sub, run["model"], X_test, y_test)
    true(acc >= case["target"], f"held-out accuracy {acc:.3f} is below {case['target']}. The reference reaches {reference:.3f}", f">= {case['target']}", acc)