import sys
from argparse import ArgumentParser
from pathlib import Path
from comet_ml import Experiment

import numpy as np
import torch
import yaml
from PIL import Image
from skimage.color import gray2rgb
from skimage.io import imread
from skimage.transform import resize
from skimage.util import img_as_ubyte
from tqdm import tqdm

sys.path.append(str(Path(__file__).resolve().parent.parent))

import climategan

GROUND_MODEL = "/miniscratch/_groups/ccai/experiments/runs/ablation-v1/out--ground"


def uint8(array):
    return array.astype(np.uint8)


def crop_and_resize(image_path, label_path):
    """
    Resizes an image so that it keeps the aspect ratio and the smallest dimensions
    is 640, then crops this resized image in its center so that the output is 640x640
    without aspect ratio distortion

    Args:
        image_path (Path or str): Path to an image
        label_path (Path or str): Path to the image's associated label

    Returns:
        tuple((np.ndarray, np.ndarray)): (new image, new label)
    """

    img = imread(image_path)
    lab = imread(label_path)

    # if img.shape[-1] == 4:
    #     img = uint8(rgba2rgb(img) * 255)

    # TODO: remove (debug)
    if img.shape[:2] != lab.shape[:2]:
        print(
            "\nWARNING: shape mismatch: im -> {}, lab -> {}".format(
                image_path.name, label_path.name
            )
        )
        # breakpoint()

    # resize keeping aspect ratio: smallest dim is 640
    h, w = img.shape[:2]
    if h < w:
        size = (640, int(640 * w / h))
    else:
        size = (int(640 * h / w), 640)

    r_img = resize(img, size, preserve_range=True, anti_aliasing=True)
    r_img = uint8(r_img)

    r_lab = resize(lab, size, preserve_range=True, anti_aliasing=False, order=0)
    r_lab = uint8(r_lab)

    # crop in the center
    H, W = r_img.shape[:2]

    top = (H - 640) // 2
    left = (W - 640) // 2

    rc_img = r_img[top : top + 640, left : left + 640, :]
    rc_lab = (
        r_lab[top : top + 640, left : left + 640, :]
        if r_lab.ndim == 3
        else r_lab[top : top + 640, left : left + 640]
    )

    return rc_img, rc_lab


def load_ground(ground_output_path, ref_image_path):
    gop = Path(ground_output_path)
    rip = Path(ref_image_path)

    ground_paths = list((gop / "eval-metrics" / "pred").glob(f"{rip.stem}.jpg")) + list(
        (gop / "eval-metrics" / "pred").glob(f"{rip.stem}.png")
    )
    if len(ground_paths) == 0:
        raise ValueError(
            f"Could not find a ground match in {str(gop)} for image {str(rip)}"
        )
    elif len(ground_paths) > 1:
        raise ValueError(
            f"Found more than 1 ground match in {str(gop)} for image {str(rip)}:"
            + f" {list(map(str, ground_paths))}"
        )
    ground_path = ground_paths[0]
    _, ground = crop_and_resize(rip, ground_path)
    ground = (ground > 0).astype(np.float32)
    return torch.from_numpy(ground).unsqueeze(0).unsqueeze(0).cuda()


def parse_args():
    parser = ArgumentParser()
    parser.add_argument("-y", "--yaml", help="Path to a list of models")
    parser.add_argument(
        "--disable_loading",
        action="store_true",
        default=False,
        help="Disable loading of existing inferences",
    )
    parser.add_argument(
        "-t", "--tags", nargs="*", help="Comet.ml tags", default=[], type=str
    )
    parser.add_argument(
        "--tasks",
        nargs="*",
        help="Comet.ml tags",
        default=["x", "d", "s", "m", "mx", "p"],
        type=str,
    )
    args = parser.parse_args()

    print("Received args:")
    print(vars(args))

    return args


def load_images_and_labels(
    path="/miniscratch/_groups/ccai/data/omnigan/masker-test-set",
):
    p = Path(path)
    ims_path = p / "imgs"
    lab_path = p / "labels"

    ims = sorted(climategan.utils.find_images(ims_path), key=lambda x: x.name)
    labs = sorted(
        climategan.utils.find_images(lab_path),
        key=lambda x: x.name.replace("_labeled.", "."),
    )

    xs = climategan.transforms.PrepareInference()(ims)
    ys = climategan.transforms.PrepareInference(is_label=True)(labs)

    return xs, ys, ims, labs


def load_inferences(inf_path, im_paths):
    try:
        assert inf_path.exists()
        assert sorted([i.stem for i in im_paths]) == sorted(
            [i.stem for i in inf_path.glob("*.pt")]
        )
        return [torch.load(str(i)) for i in tqdm(list(inf_path.glob("*.pt")))]
    except Exception as e:
        print()
        print(e)
        print("Aborting Loading")
        print()
        return None


def get_or_load_inferences(
    m_path, device, xs, is_ground, im_paths, ground_model, try_load=True
):
    inf_path = Path(m_path) / "inferences"
    if try_load:
        print("Trying to load existing inferences:")
        outputs = load_inferences(inf_path, im_paths)
        if outputs is not None:
            print("Successfully loaded existing inferences")
            return outputs

    trainer = climategan.trainer.Trainer.resume_from_path(
        m_path if not is_ground else ground_model,
        inference=True,
        new_exp=None,
        device=device,
    )

    inf_path.mkdir(exist_ok=True)
    outputs = []
    for i, x in enumerate(tqdm(xs)):
        x = x.to(trainer.device)
        if not is_ground:
            out = trainer.G.decode(x=x)
        else:
            out = {"m": load_ground(GROUND_MODEL, im_paths[i])}
        out["p"] = trainer.G.paint(out["m"] > 0.5, x)
        out["x"] = x
        inference = {k: v.cpu() for k, v in out.items()}
        outputs.append(inference)
        torch.save(inference, inf_path / f"{im_paths[i].stem}.pt")
    print()

    return outputs


def numpify(outputs):
    nps = []
    print("Numpifying...")
    for o in tqdm(outputs):
        x = (o["x"][0].permute(1, 2, 0).numpy() + 1) / 2
        m = o["m"]
        m = (m[0, 0, :, :].numpy() > 0.5).astype(np.uint8)
        p = (o["p"][0].permute(1, 2, 0).numpy() + 1) / 2
        data = {"m": m, "p": p, "x": x}
        if "s" in o:
            s = climategan.data.decode_segmap_merged_labels(o["s"], "r", False) / 255.0
            data["s"] = s[0].permute(1, 2, 0).numpy()
        if "d" in o:
            d = climategan.tutils.normalize_tensor(o["d"]).squeeze().numpy()
            data["d"] = d
        nps.append({k: img_as_ubyte(v) for k, v in data.items()})
    return nps


def concat_npy_for_model(data, tasks):
    assert "m" in data
    assert "x" in data
    assert "p" in data

    x = mask = depth = seg = painted = masked = None

    x = data["x"]
    painted = data["p"]
    mask = (gray2rgb(data["m"]) * 255).astype(np.uint8)
    painted = data["p"]
    masked = (1 - gray2rgb(data["m"])) * x

    concats = []

    if "d" in data:
        depth = img_as_ubyte(
            gray2rgb(
                resize(data["d"], data["x"].shape[:2], anti_aliasing=True, order=1)
            )
        )
    else:
        depth = np.ones_like(data["x"]) * 255

    if "s" in data:
        seg = img_as_ubyte(
            resize(data["s"], data["x"].shape[:2], anti_aliasing=False, order=0)
        )
    else:
        seg = np.ones_like(data["x"]) * 255

    for t in tasks:
        if t == "x":
            concats.append(x)
        if t == "m":
            concats.append(mask)
        elif t == "mx":
            concats.append(masked)
        elif t == "d":
            concats.append(depth)
        elif t == "s":
            concats.append(seg)
        elif t == "p":
            concats.append(painted)

    row = np.concatenate(concats, axis=1)

    return row


if __name__ == "__main__":
    args = parse_args()

    with open(args.yaml, "r") as f:
        maskers = yaml.safe_load(f)
    if "models" in maskers:
        maskers = maskers["models"]

    load = not args.disable_loading
    tags = args.tags
    tasks = args.tasks

    ground_model = None
    for m in maskers:
        if "ground" not in maskers:
            ground_model = m
            break
    if ground_model is None:
        raise ValueError("Could not find a non-ground model to get a painter")

    device = torch.device("cuda:0")
    torch.set_grad_enabled(False)

    xs, ys, im_paths, lab_paths = load_images_and_labels()

    np_outs = {}
    names = []

    for m_path in maskers:

        opt_path = Path(m_path) / "opts.yaml"
        with opt_path.open("r") as f:
            opt = yaml.safe_load(f)

        name = (
            ", ".join(
                [
                    t
                    for t in sorted(opt["comet"]["tags"])
                    if "branch" not in t and "ablation" not in t and "trash" not in t
                ]
            )
            if "--ground" not in m_path
            else "ground"
        )
        names.append(name)

        is_ground = name == "ground"

        print("#" * 100)
        print("\n>>> Processing", name)
        print()

        outputs = get_or_load_inferences(
            m_path, device, xs, is_ground, im_paths, ground_model, load
        )
        nps = numpify(outputs)

        np_outs[name] = nps

    exp = Experiment(project_name="climategan-inferences", display_summary_level=0)
    exp.log_parameter("names", names)
    exp.add_tags(tags)

    for i in tqdm(range(len(xs))):
        all_models_for_image = []
        for name in names:
            xpmds = concat_npy_for_model(np_outs[name][i], tasks)
            all_models_for_image.append(xpmds)
        full_im = np.concatenate(all_models_for_image, axis=0)
        pil_im = Image.fromarray(full_im)
        exp.log_image(pil_im, name=im_paths[i].stem.replace(".", "_"), step=i)