import argparse
import json
import os
import random
from copy import deepcopy

import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from timm.models.layers import DropPath, trunc_normal_
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.transforms import functional


# ========= Model =========

# copy from https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py


class LayerNorm(nn.Module):
    """LayerNorm that supports two data formats: channels_last (default) or channels_first.
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
    with shape (batch_size, channels, height, width).
    """

    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError
        self.normalized_shape = (normalized_shape,)

    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(
                x, self.normalized_shape, self.weight, self.bias, self.eps
            )
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class GRN(nn.Module):
    """GRN (Global Response Normalization) layer"""

    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x


class Block(nn.Module):
    """ConvNeXtV2 Block.

    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
    """

    def __init__(self, dim, drop_path=0.0):
        super().__init__()
        self.dwconv = nn.Conv2d(
            dim, dim, kernel_size=7, padding=3, groups=dim
        )  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(
            dim, 4 * dim
        )  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.grn = GRN(4 * dim)
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.grn(x)
        x = self.pwconv2(x)
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x


class ConvNeXtV2(nn.Module):
    """ConvNeXt V2

    Args:
        in_chans (int): Number of input image channels. Default: 3
        num_classes (int): Number of classes for classification head. Default: 1000
        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
        dims (int[]): Feature dimension at each stage. Default: [96, 192, 384, 768]
        drop_path_rate (float): Stochastic depth rate. Default: 0.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """

    def __init__(
            self,
            in_chans=3,
            num_classes=1000,
            depths=[3, 3, 9, 3],
            dims=[96, 192, 384, 768],
            drop_path_rate=0.0,
            head_init_scale=1.0,
    ):
        super().__init__()
        self.depths = depths
        self.downsample_layers = (
            nn.ModuleList()
        )  # stem and 3 intermediate downsampling conv layers
        stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
            LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
        )
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

        self.stages = (
            nn.ModuleList()
        )  # 4 feature resolution stages, each consisting of multiple residual blocks
        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[
                    Block(dim=dims[i], drop_path=dp_rates[cur + j])
                    for j in range(depths[i])
                ]
            )
            self.stages.append(stage)
            cur += depths[i]

        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
        self.head = nn.Linear(dims[-1], num_classes)

        self.apply(self._init_weights)
        self.head.weight.data.mul_(head_init_scale)
        self.head.bias.data.mul_(head_init_scale)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=0.02)
            nn.init.constant_(m.bias, 0)

    def forward_features(self, x):
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)
        return self.norm(
            x.mean([-2, -1])
        )  # global average pooling, (N, C, H, W) -> (N, C)

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x


model_cfgs = {
    "atto": [[2, 2, 6, 2], [40, 80, 160, 320]],
    "femto": [[2, 2, 6, 2], [48, 96, 192, 384]],
    "pico": [[2, 2, 6, 2], [64, 128, 256, 512]],
    "nano": [[2, 2, 8, 2], [80, 160, 320, 640]],
    "tiny": [[3, 3, 9, 3], [96, 192, 384, 768]],
    "base": [[3, 3, 27, 3], [128, 256, 512, 1024]],
    "large": [[3, 3, 27, 3], [192, 384, 768, 1536]],
    "huge": [[3, 3, 27, 3], [352, 704, 1408, 2816]],
}


def convnextv2(cfg_name, **kwargs):
    cfg = model_cfgs[cfg_name]
    model = ConvNeXtV2(depths=cfg[0], dims=cfg[1], **kwargs)
    return model


# ========= Dataset =========

EXTENSION = [".png", ".jpg", ".jpeg"]


def file_ext(fname):
    return os.path.splitext(fname)[1].lower()


def rescale_pad(image, output_size, random_pad=False):
    h, w = image.shape[-2:]
    if h != output_size or w != output_size:
        r = min(output_size / h, output_size / w)
        new_h, new_w = int(h * r), int(w * r)
        ph = output_size - new_h
        pw = output_size - new_w
        image = transforms.functional.resize(image, [new_h, new_w])
        image = transforms.functional.pad(
            image, [pw // 2, ph // 2, pw // 2 + pw % 2, ph // 2 + ph % 2], random.uniform(0, 1) if random_pad else 0
        )
    return image


def random_crop(image, min_rate=0.8):
    h, w = image.shape[-2:]
    new_h, new_w = int(h * random.uniform(min_rate, 1)), int(w * random.uniform(min_rate, 1))
    top = np.random.randint(0, h - new_h)
    left = np.random.randint(0, w - new_w)
    image = image[:, top: top + new_h, left: left + new_w]
    return image


class AnimeAestheticDataset(Dataset):
    def __init__(self, path, img_size, xflip=True):
        all_files = {
            os.path.relpath(os.path.join(root, fname), path)
            for root, _dirs, files in os.walk(path)
            for fname in files
        }
        all_images = sorted(
            fname for fname in all_files if file_ext(fname) in EXTENSION
        )
        with open(os.path.join(path, "label.json"), "r", encoding="utf8") as f:
            labels = json.load(f)
        image_list = []
        label_list = []
        for fname in all_images:
            if fname not in labels:
                continue
            image_list.append(fname)
            label_list.append(labels[fname])
        self.path = path
        self.img_size = img_size
        self.xflip = xflip
        self.image_list = image_list
        self.label_list = label_list

    def __len__(self):
        length = len(self.image_list)
        if self.xflip:
            length *= 2
        return length

    def __getitem__(self, index):
        real_len = len(self.image_list)
        fname = self.image_list[index % real_len]
        label = self.label_list[index % real_len]
        image = Image.open(os.path.join(self.path, fname)).convert("RGB")
        image = transforms.functional.to_tensor(image)
        image = random_crop(image, 0.8)
        image = rescale_pad(image, self.img_size, True)
        if index // real_len != 0:
            image = transforms.functional.hflip(image)
        label = torch.tensor([label], dtype=torch.float32)
        return image, label


# ========= Train =========


class AnimeAesthetic(pl.LightningModule):
    def __init__(self, cfg: str, drop_path_rate=0.0, ema_decay=0):
        super().__init__()
        self.net = convnextv2(cfg, in_chans=3, num_classes=1, drop_path_rate=drop_path_rate)
        self.ema_decay = ema_decay
        self.ema = None
        if ema_decay > 0:
            self.ema = deepcopy(self.net)
            self.ema.requires_grad_(False)

    def configure_optimizers(self):
        optimizer = optim.Adam(
            self.net.parameters(),
            lr=0.001,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=0,
        )
        return optimizer

    def forward(self, x, use_ema=False):
        x = (x - 0.5) / 0.5
        net = self.ema if use_ema else self.net
        return net(x)

    def training_step(self, batch, batch_idx):
        images, labels = batch
        loss = F.mse_loss(self.forward(images, False), labels)
        self.log_dict({"train/loss": loss})
        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        mae = F.l1_loss(self.forward(images, False), labels)
        logs = {"val/mae": mae}
        if self.ema is not None:
            mae_ema = F.l1_loss(self.forward(images, True), labels)
            logs["val/mae_ema"] = mae_ema
        self.log_dict(logs, sync_dist=True)

    def on_train_batch_end(self, outputs, batch, batch_idx):
        if self.ema is not None:
            with torch.no_grad():
                for ema_v, model_v in zip(
                        self.ema.state_dict().values(), self.net.state_dict().values()
                ):
                    ema_v.copy_(
                        self.ema_decay * ema_v + (1.0 - self.ema_decay) * model_v
                    )


def main(opt):
    if not os.path.exists("lightning_logs"):
        os.mkdir("lightning_logs")
    torch.manual_seed(0)
    np.random.seed(0)
    print("---load dataset---")
    full_dataset = AnimeAestheticDataset(opt.data, opt.img_size)
    full_dataset_len = len(full_dataset)
    train_dataset_len = int(full_dataset_len * opt.data_split)
    val_dataset_len = full_dataset_len - train_dataset_len
    train_dataset, val_dataset = torch.utils.data.random_split(
        full_dataset, [train_dataset_len, val_dataset_len]
    )
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=opt.batch_size_train,
        shuffle=True,
        persistent_workers=True,
        num_workers=opt.workers_train,
        pin_memory=True,
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=opt.batch_size_val,
        shuffle=False,
        persistent_workers=True,
        num_workers=opt.workers_val,
        pin_memory=True,
    )
    print(f"train: {len(train_dataset)}")
    print(f"val: {len(val_dataset)}")
    print("---define model---")
    if opt.resume != "":
        anime_aesthetic = AnimeAesthetic.load_from_checkpoint(
            opt.resume, cfg=opt.cfg, drop_path_rate=opt.drop_path, ema_decay=opt.ema_decay
        )
    else:
        anime_aesthetic = AnimeAesthetic(cfg=opt.cfg, drop_path_rate=opt.drop_path, ema_decay=opt.ema_decay)

    print("---start train---")

    checkpoint_callback = ModelCheckpoint(
        monitor="val/mae",
        mode="min",
        save_top_k=1,
        save_last=True,
        auto_insert_metric_name=False,
        filename="epoch={epoch},mae={val/mae:.4f}",
    )
    callbacks = [checkpoint_callback]
    if opt.ema_decay > 0:
        checkpoint_ema_callback = ModelCheckpoint(
            monitor="val/mae_ema",
            mode="min",
            save_top_k=1,
            save_last=False,
            auto_insert_metric_name=False,
            filename="epoch={epoch},mae-ema={val/mae_ema:.4f}",
        )
        callbacks.append(checkpoint_ema_callback)
    trainer = Trainer(
        precision=32 if opt.fp32 else 16,
        accelerator=opt.accelerator,
        devices=opt.devices,
        max_epochs=opt.epoch,
        benchmark=opt.benchmark,
        accumulate_grad_batches=opt.acc_step,
        val_check_interval=opt.val_epoch,
        log_every_n_steps=opt.log_step,
        strategy="ddp_find_unused_parameters_false" if opt.devices > 1 else None,
        callbacks=callbacks,
    )
    trainer.fit(anime_aesthetic, train_dataloader, val_dataloader)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # model args
    parser.add_argument(
        "--cfg",
        type=str,
        default="tiny",
        choices=list(model_cfgs.keys()),
        help="model configure",
    )
    parser.add_argument(
        "--resume", type=str, default="", help="resume training from ckpt"
    )
    parser.add_argument(
        "--img-size",
        type=int,
        default=768,
        help="image size for training and validation",
    )

    # dataset args
    parser.add_argument(
        "--data", type=str, default="./data", help="dataset path"
    )
    parser.add_argument(
        "--data-split",
        type=float,
        default=0.9995,
        help="split rate for training and validation",
    )

    # training args
    parser.add_argument("--epoch", type=int, default=100, help="epoch num")
    parser.add_argument(
        "--batch-size-train", type=int, default=16, help="batch size for training"
    )
    parser.add_argument(
        "--batch-size-val", type=int, default=2, help="batch size for val"
    )
    parser.add_argument(
        "--workers-train",
        type=int,
        default=4,
        help="workers num for training dataloader",
    )
    parser.add_argument(
        "--workers-val",
        type=int,
        default=4,
        help="workers num for validation dataloader",
    )
    parser.add_argument(
        "--acc-step", type=int, default=8, help="gradient accumulation step"
    )
    parser.add_argument(
        "--drop-path", type=float, default=0.1, help="Drop path rate"
    )
    parser.add_argument(
        "--ema-decay", type=float, default=0.9999, help="use ema if ema-decay > 0"
    )
    parser.add_argument(
        "--accelerator",
        type=str,
        default="gpu",
        choices=["cpu", "gpu", "tpu", "ipu", "hpu", "auto"],
        help="accelerator",
    )
    parser.add_argument("--devices", type=int, default=4, help="devices num")
    parser.add_argument(
        "--fp32", action="store_true", default=False, help="disable mix precision"
    )
    parser.add_argument(
        "--benchmark", action="store_true", default=True, help="enable cudnn benchmark"
    )
    parser.add_argument(
        "--log-step", type=int, default=2, help="log training loss every n steps"
    )
    parser.add_argument(
        "--val-epoch", type=int, default=0.1, help="valid and save every n epoch"
    )

    opt = parser.parse_args()
    print(opt)

    main(opt)