import time
from typing import Callable, Dict, List, Tuple

import matplotlib.pyplot as plt
import numpy as np
import torch as tc
import torch.nn as nn
from PIL import Image
from torch.utils.data import Dataset

DEVICE = tc.device("cuda" if tc.cuda.is_available() else "cpu")


def model_size(m: nn.Module):
    a = sum(p.numel() for p in m.parameters())
    b = sum(p.numel() for p in m.parameters() if p.requires_grad)
    return f"trainable: {b}/{a}"


class WindowDataset(Dataset):
    def __init__(s, x, W):
        s.x, s.W = x, W

    def __len__(s):
        return len(s.x) - s.W + 1

    def __getitem__(s, idx):
        return s.x[idx : idx + s.W]


# ==============================

D_TYPE = Dict[str, np.ndarray]
D2_TYPE = Dict[str, D_TYPE]
T2_TYPE = Dict[str, Dict[str, tc.Tensor]]
LOSS_FN_TYPE = Callable[[T2_TYPE, str], Tuple[tc.Tensor, Dict]]


def train_model(
    net: nn.Module,
    xy_train: D2_TYPE,
    xy_test: D2_TYPE,
    loss_fn: LOSS_FN_TYPE,
    id="train_model",
    lr=1e-3,
    n_epoch=10000,
    f_test=10,
    f_plot=20,
):
    net.to(DEVICE)
    xy_train, xy_test = tensor(xy_train), tensor(xy_test)
    opt = tc.optim.AdamW(net.parameters(), lr=lr)
    records, gif = [], GIFMaker()

    for e in range(n_epoch):
        loss_test, info_test = np.nan, {}
        if e % f_test == 0:
            net.eval()
            with tc.no_grad():
                loss_test, info_test, plots_test = loss_fn(xy_test, "test")
            net.train()
        opt.zero_grad()
        loss_train, info_train, plots_train = loss_fn(xy_train, "train")
        loss_train.backward()
        opt.step()
        rec = {"a_loss_train": loss_train, "a_loss_test": loss_test}
        records.append({**rec, **info_train, **info_test})
        if e % f_plot == 0:
            plots = transpose_records(records)
            plots = {**plots, **plots_train, **plots_test}
            plot_general(dict(sorted(plots.items())), f"{id}.png")
            gif.add(f"{id}.png")
    gif.save(id)


def mod_keys(dic: Dict, post):
    return {f"{k}_{post}": v for k, v in dic.items()}


def slice_xy(xy: D2_TYPE, r1, r2, num=None):
    xy2 = {}
    for sym, d in xy.items():
        x, y = d["x"], d["y"]
        n = len(x)
        i1, i2 = int(r1 * n), int(r2 * n)
        step = max(1, (i2 - i1) // num) if num else 1
        x, y = x[i1:i2:step], y[i1:i2:step]
        xy2[sym] = dict(x=x, y=y)
    print(f"slice_xy: n: {n}, i1: {i1}, i2: {i2}, step: {step}, n2: {len(x)}")
    return xy2


def concat_xy(xy: D2_TYPE):
    X = np.concat([d["x"] for d in xy.values()], axis=0)
    Y = np.concat([d["y"] for d in xy.values()], axis=0)
    return {"ALL_SYMBOLS": dict(x=X, y=Y)}


def timer(func):
    def func2(*args, **kwargs):
        t1 = time.time()
        res = func(*args, **kwargs)
        print(f"{func.__name__} took {time.time() - t1:.3f} seconds")
        return res

    return func2


def tensor(x):
    if isinstance(x, np.ndarray):
        return tc.from_numpy(x.copy()).float().to(DEVICE)
    if isinstance(x, dict):
        return {k: tensor(v) for k, v in x.items()}
    if isinstance(x, list):
        return [tensor(v) for v in x]
    return x


def to_np(x):
    if isinstance(x, tc.Tensor):
        return x.detach().cpu().numpy()
    if isinstance(x, dict):
        return {k: to_np(v) for k, v in x.items()}
    if isinstance(x, list):
        return np.array([to_np(v) for v in x])
    return x


def shape(x):
    if isinstance(x, (np.ndarray, tc.Tensor)):
        return x.shape
    if isinstance(x, dict):
        return {k: shape(v) for k, v in x.items()}
    if isinstance(x, list):
        return {shape(v) for v in x}


def l2_norm(x: nn.Module):
    return sum(p.norm(2) for p in x.parameters())


def add_l2_loss(loss: tc.Tensor, net: nn.Module, ratio):
    l2 = l2_norm(net)
    loss += l2 / l2.item() * loss.abs().item() * ratio
    return loss, l2


# ==================================


def plot_general(plots: D_TYPE, id, C=2, bins=100):
    R = int(np.ceil(len(plots) / C))
    plt.figure(figsize=(4 * C, 3 * R))
    i = 0
    for k, v in plots.items():
        v = to_np(v)
        i += 1
        plt.subplot(R, C, i)
        plt.title(k)
        if "_hist" in k:
            v = v.flatten()
            if len(v):
                mean, std = v.mean(), v.std() * 3
                plt.hist(v, bins=bins, range=[mean - std, mean + std])
        else:
            x = np.arange(len(v))
            if np.any(np.isnan(v)):
                plt.scatter(x, v, s=3)
            else:
                plt.plot(x, v)
    plt.tight_layout()
    plt.savefig(id)
    plt.close()


def transpose_records(records: List[Dict]):
    return {k: np.array([to_np(r.get(k, np.nan)) for r in records]) for k in records[0]}


def plot_records(records: List[Dict], id, C=1):
    dic = transpose_records(records)
    R = int(np.ceil(len(dic) / C))
    plt.figure(figsize=(4 * C, 3 * R))
    i = 0
    for k, v in dic.items():
        i += 1
        plt.subplot(R, C, i)
        plt.title(k)
        x = np.arange(len(v))
        if np.any(np.isnan(v)):
            plt.scatter(x, v, s=3)
        else:
            plt.plot(x, v)
    plt.tight_layout()
    plt.savefig(id)
    plt.close()


def plot_xy(xy: D2_TYPE, id="xy"):
    R, C = 5, 2
    plt.figure(figsize=(4 * C, 3 * R))
    i = 0
    for sym, d in xy.items():
        i += 1
        x, y = d["x"], d["y"]
        plt.subplot(R, C, i)
        plt.title(sym)
        for k in range(x.shape[1]):
            plt.plot(x[:, k], label=f"x[:, {k}]")
        plt.legend()
        i += 1
        plt.subplot(R, C, i)
        plt.plot(y)
        if i == R * C:
            break
    plt.tight_layout()
    plt.savefig(id)
    plt.close()


# =====================


class GIFMaker:
    def __init__(s):
        s.frames: List[Image.Image] = []

    def add(s, path):
        s.frames.append(Image.open(path).copy())

    def save(s, id, fps=10):
        s.frames[0].save(
            f"{id}.gif",
            format="GIF",
            append_images=s.frames[1:],
            save_all=True,
            duration=int(1000 / fps),
            loop=0,
            optimize=True,
        )
