from functools import cache
from pickle import load

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import PIL.Image as Image
import torch

from msma import ScoreFlow, config_presets


@cache
def load_model(modeldir, preset="edm2-img64-s-fid", device="cpu", outdir=None):
    model = ScoreFlow(preset, device=device)
    model.flow.load_state_dict(torch.load(f"{modeldir}/{preset}/flow.pt"))
    return model


@cache
def load_reference_scores(model_dir):
    with np.load(f"{model_dir}/refscores.npz", "rb") as f:
        ref_nll = f["arr_0"]
    return ref_nll


def compute_gmm_likelihood(x_score, model_dir):
    with open(f"{model_dir}/gmm.pkl", "rb") as f:
        clf = load(f)
        nll = -clf.score(x_score)

    ref_nll = load_reference_scores(model_dir)
    percentile = (ref_nll < nll).mean() * 100

    return nll, percentile, ref_nll


def plot_against_reference(nll, ref_nll):
    fig, ax = plt.subplots()
    ax.hist(ref_nll, label="Reference Scores")
    ax.axvline(nll, label="Image Score", c="red", ls="--")
    plt.legend()
    fig.tight_layout()
    return fig


def plot_heatmap(img: Image, heatmap: np.array):
    fig, ax = plt.subplots()
    cmap = plt.get_cmap("gist_heat")
    h = -heatmap[0, 0].copy()
    qmin, qmax = np.quantile(h, 0.8), np.quantile(h, 0.999)
    h = np.clip(h, a_min=qmin, a_max=qmax)
    h = (h - h.min()) / (h.max() - h.min())
    h = cmap(h, bytes=True)[:, :, :3]
    h = Image.fromarray(h).resize(img.size, resample=Image.Resampling.BILINEAR)
    im = Image.blend(img, h, alpha=0.6)
    # im = ax.imshow(np.array(im))
    # # fig.colorbar(im)
    # # plt.grid(False)
    # # plt.axis("off")
    # fig.tight_layout()
    return im


def run_inference(input_img, preset="edm2-img64-s-fid",  device="cuda"):

    # img = center_crop_imagenet(64, img)
    input_img = input_img.resize(size=(64, 64), resample=Image.Resampling.LANCZOS)

    with torch.inference_mode():
        img = np.array(input_img)
        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
        img = img.float().to(device)
        model = load_model(modeldir="models", preset=preset, device=device)
        img_likelihood = model(img).cpu().numpy()
        # img_likelihood = model.scorenet(img).square().sum(1).sum(1).contiguous().float().cpu().unsqueeze(1).numpy()
        # print(img_likelihood.shape, img_likelihood.dtype)
        img = torch.nn.functional.interpolate(img, size=64, mode="bilinear")
        x = model.scorenet(img)
        x = x.square().sum(dim=(2, 3, 4)) ** 0.5
        nll, pct, ref_nll = compute_gmm_likelihood(
            x.cpu(), model_dir=f"models/{preset}"
        )

    outstr = f"Anomaly score: {nll:.3f} / {pct:.2f} percentile"
    histplot = plot_against_reference(nll, ref_nll)
    heatmapplot = plot_heatmap(input_img, img_likelihood)

    return outstr, heatmapplot, histplot


demo = gr.Interface(
    fn=run_inference,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Dropdown(choices=config_presets.keys(), label="Score Model"),
    ],
    outputs=[
        "text",
        gr.Image(label="Anomaly Heatmap", min_width=64),
        gr.Plot(label="Comparing to Imagenette"),
    ],

    examples=[
        ['goldfish.JPEG', "edm2-img64-s-fid"]
    ]
)

if __name__ == "__main__":
    demo.launch()