import torch
import torch.distributed as dist
from torchvision import transforms as tvtrans
import os
import os.path as osp
import time
import timeit
import copy
import json
import pickle
import PIL.Image
import numpy as np
from datetime import datetime
from easydict import EasyDict as edict
from collections import OrderedDict

from lib.cfg_holder import cfg_unique_holder as cfguh
from lib.data_factory import get_dataset, get_sampler, collate
from lib.model_zoo import \
    get_model, get_optimizer, get_scheduler
from lib.log_service import print_log

from ..utils import train as train_base
from ..utils import eval as eval_base
from ..utils import train_stage as tsbase
from ..utils import eval_stage as esbase
from .. import sync

###############
# some helper #
###############

def atomic_save(cfg, net, opt, step, path):
    if isinstance(net, (torch.nn.DataParallel,
                        torch.nn.parallel.DistributedDataParallel)):
        netm = net.module
    else:
        netm = net
    sd = netm.state_dict()
    slimmed_sd = [(ki, vi) for ki, vi in sd.items()
        if ki.find('first_stage_model')!=0 and ki.find('cond_stage_model')!=0]

    checkpoint = {
        "config" : cfg,
        "state_dict" : OrderedDict(slimmed_sd),
        "step" : step}
    if opt is not None:
        checkpoint['optimizer_states'] = opt.state_dict()
    import io
    import fsspec
    bytesbuffer = io.BytesIO()
    torch.save(checkpoint, bytesbuffer)
    with fsspec.open(path, "wb") as f:
        f.write(bytesbuffer.getvalue())

def load_state_dict(net, cfg):
    pretrained_pth_full  = cfg.get('pretrained_pth_full' , None)
    pretrained_ckpt_full = cfg.get('pretrained_ckpt_full', None)
    pretrained_pth       = cfg.get('pretrained_pth'      , None)
    pretrained_ckpt      = cfg.get('pretrained_ckpt'     , None)
    pretrained_pth_dm    = cfg.get('pretrained_pth_dm'   , None)
    pretrained_pth_ema   = cfg.get('pretrained_pth_ema'  , None)
    strict_sd = cfg.get('strict_sd', False)
    errmsg = "Overlapped model state_dict! This is undesired behavior!"

    if pretrained_pth_full is not None or pretrained_ckpt_full is not None:
        assert (pretrained_pth is None) and \
               (pretrained_ckpt is None) and \
               (pretrained_pth_dm is None) and \
               (pretrained_pth_ema is None), errmsg            
        if pretrained_pth_full is not None:
            target_file = pretrained_pth_full
            sd = torch.load(target_file, map_location='cpu')
            assert pretrained_ckpt is None, errmsg
        else:
            target_file = pretrained_ckpt_full
            sd = torch.load(target_file, map_location='cpu')['state_dict']
        print_log('Load full model from [{}] strict [{}].'.format(
            target_file, strict_sd))
        net.load_state_dict(sd, strict=strict_sd)

    if pretrained_pth is not None or pretrained_ckpt is not None:
        assert (pretrained_ckpt_full is None) and \
               (pretrained_pth_full is None) and \
               (pretrained_pth_dm is None) and \
               (pretrained_pth_ema is None), errmsg
        if pretrained_pth is not None:
            target_file = pretrained_pth
            sd = torch.load(target_file, map_location='cpu')
            assert pretrained_ckpt is None, errmsg
        else:
            target_file = pretrained_ckpt
            sd = torch.load(target_file, map_location='cpu')['state_dict']
        print_log('Load model from [{}] strict [{}].'.format(
            target_file, strict_sd))
        sd_extra = [(ki, vi) for ki, vi in net.state_dict().items() \
            if ki.find('first_stage_model')==0 or ki.find('cond_stage_model')==0]
        sd.update(OrderedDict(sd_extra))
        net.load_state_dict(sd, strict=strict_sd)

    if pretrained_pth_dm is not None:
        assert (pretrained_ckpt_full is None) and \
               (pretrained_pth_full is None) and \
               (pretrained_pth is None) and \
               (pretrained_ckpt is None), errmsg
        print_log('Load diffusion model from [{}] strict [{}].'.format(
            pretrained_pth_dm, strict_sd))
        sd = torch.load(pretrained_pth_dm, map_location='cpu')
        net.model.diffusion_model.load_state_dict(sd, strict=strict_sd)

    if pretrained_pth_ema is not None:
        assert (pretrained_ckpt_full is None) and \
               (pretrained_pth_full is None) and \
               (pretrained_pth is None) and \
               (pretrained_ckpt is None), errmsg
        print_log('Load unet ema model from [{}] strict [{}].'.format(
            pretrained_pth_ema, strict_sd))
        sd = torch.load(pretrained_pth_ema, map_location='cpu')
        net.model_ema.load_state_dict(sd, strict=strict_sd)

def auto_merge_imlist(imlist, max=64):
    imlist = imlist[0:max]
    h, w = imlist[0].shape[0:2]
    num_images = len(imlist)
    num_row = int(np.sqrt(num_images))
    num_col = num_images//num_row + 1 if num_images%num_row!=0 else num_images//num_row
    canvas = np.zeros([num_row*h, num_col*w, 3], dtype=np.uint8)
    for idx, im in enumerate(imlist):
        hi = (idx // num_col) * h
        wi = (idx % num_col) * w
        canvas[hi:hi+h, wi:wi+w, :] = im
    return canvas

def latent2im(net, latent):
    single_input = len(latent.shape) == 3
    if single_input:
        latent = latent[None]
    im = net.decode_image(latent.to(net.device))
    im = torch.clamp((im+1.0)/2.0, min=0.0, max=1.0)
    im = [tvtrans.ToPILImage()(i) for i in im]
    if single_input:
        im = im[0]
    return im

def im2latent(net, im):
    single_input = not isinstance(im, list)
    if single_input:
        im = [im]
    im = torch.stack([tvtrans.ToTensor()(i) for i in im], dim=0)
    im = (im*2-1).to(net.device)
    z = net.encode_image(im)
    if single_input:
        z = z[0]
    return z

class color_adjust(object):
    def __init__(self, ref_from, ref_to):
        x0, m0, std0 = self.get_data_and_stat(ref_from)
        x1, m1, std1 = self.get_data_and_stat(ref_to)
        self.ref_from_stat = (m0, std0)
        self.ref_to_stat   = (m1, std1)
        self.ref_from = self.preprocess(x0).reshape(-1, 3)
        self.ref_to = x1.reshape(-1, 3)

    def get_data_and_stat(self, x):
        if isinstance(x, str):
            x = np.array(PIL.Image.open(x))
        elif isinstance(x, PIL.Image.Image):
            x = np.array(x)
        elif isinstance(x, torch.Tensor):
            x = torch.clamp(x, min=0.0, max=1.0)
            x = np.array(tvtrans.ToPILImage()(x))
        elif isinstance(x, np.ndarray):
            pass
        else:
            raise ValueError
        x = x.astype(float)
        m = np.reshape(x, (-1, 3)).mean(0)
        s = np.reshape(x, (-1, 3)).std(0)
        return x, m, s

    def preprocess(self, x):
        m0, s0 = self.ref_from_stat
        m1, s1 = self.ref_to_stat
        y = ((x-m0)/s0)*s1 + m1
        return y

    def __call__(self, xin, keep=0, simple=False):
        xin, _, _ = self.get_data_and_stat(xin)
        x = self.preprocess(xin)
        if simple: 
            y = (x*(1-keep) + xin*keep)
            y = np.clip(y, 0, 255).astype(np.uint8)
            return y

        h, w = x.shape[:2]
        x = x.reshape(-1, 3)
        y = []
        for chi in range(3):
            yi = self.pdf_transfer_1d(self.ref_from[:, chi], self.ref_to[:, chi], x[:, chi])
            y.append(yi)

        y = np.stack(y, axis=1)
        y = y.reshape(h, w, 3)
        y = (y.astype(float)*(1-keep) + xin.astype(float)*keep)
        y = np.clip(y, 0, 255).astype(np.uint8)
        return y

    def pdf_transfer_1d(self, arr_fo, arr_to, arr_in, n=600):
        arr = np.concatenate((arr_fo, arr_to))
        min_v = arr.min() - 1e-6
        max_v = arr.max() + 1e-6
        min_vto = arr_to.min() - 1e-6
        max_vto = arr_to.max() + 1e-6
        xs = np.array(
            [min_v + (max_v - min_v) * i / n for i in range(n + 1)])
        hist_fo, _ = np.histogram(arr_fo, xs)
        hist_to, _ = np.histogram(arr_to, xs)
        xs = xs[:-1]
        # compute probability distribution
        cum_fo = np.cumsum(hist_fo)
        cum_to = np.cumsum(hist_to)
        d_fo = cum_fo / cum_fo[-1]
        d_to = cum_to / cum_to[-1]
        # transfer
        t_d = np.interp(d_fo, d_to, xs)
        t_d[d_fo <= d_to[ 0]] = min_vto
        t_d[d_fo >= d_to[-1]] = max_vto
        arr_out = np.interp(arr_in, xs, t_d)
        return arr_out

########
# main #
########

class eval(eval_base):
    def prepare_model(self):
        cfg = cfguh().cfg
        net = get_model()(cfg.model)
        if cfg.env.cuda:
            net.to(self.local_rank)
            load_state_dict(net, cfg.eval) #<--- added
            net = torch.nn.parallel.DistributedDataParallel(
                net, device_ids=[self.local_rank], 
                find_unused_parameters=True)
        net.eval()
        return {'net' : net,}

class eval_stage(esbase):
    """
    This is eval stage that can check comprehensive results
    """
    def __init__(self):
        from ..model_zoo.ddim import DDIMSampler
        self.sampler = DDIMSampler

    def get_net(self, paras):
        return paras['net']

    def get_image_path(self):
        if 'train' in cfguh().cfg:
            log_dir = cfguh().cfg.train.log_dir
        else:
            log_dir = cfguh().cfg.eval.log_dir
        return os.path.join(log_dir, "udemo")

    @torch.no_grad()
    def sample(self, net, sampler, prompt, output_dim, scale, n_samples, ddim_steps, ddim_eta):
        h, w = output_dim
        uc = None
        if scale != 1.0:
            uc = net.get_learned_conditioning(n_samples * [""])
        c = net.get_learned_conditioning(n_samples * [prompt])
        shape = [4, h//8, w//8]
        rv = sampler.sample(
            S=ddim_steps,
            conditioning=c,
            batch_size=n_samples,
            shape=shape,
            verbose=False,
            unconditional_guidance_scale=scale,
            unconditional_conditioning=uc,
            eta=ddim_eta)
        return rv

    def save_images(self, pil_list, name, path, suffix=''):
        canvas = auto_merge_imlist([np.array(i) for i in pil_list])
        image_name = '{}{}.png'.format(name, suffix)
        PIL.Image.fromarray(canvas).save(osp.join(path, image_name))

    def __call__(self, **paras):
        cfg = cfguh().cfg
        cfgv = cfg.eval

        net = paras['net']
        eval_cnt = paras.get('eval_cnt', None)
        fix_seed = cfgv.get('fix_seed', False)

        LRANK = sync.get_rank('local')
        LWSIZE = sync.get_world_size('local')

        image_path = self.get_image_path()
        self.create_dir(image_path)
        eval_cnt = paras.get('eval_cnt', None)
        suffix='' if eval_cnt is None else '_itern'+str(eval_cnt)

        if isinstance(net, (torch.nn.DataParallel,
                            torch.nn.parallel.DistributedDataParallel)):
            netm = net.module
        else:
            netm = net

        with_ema = getattr(netm, 'model_ema', None) is not None
        sampler = self.sampler(netm)
        setattr(netm, 'device', LRANK) # Trick

        replicate = cfgv.get('replicate', 1)
        conditioning = cfgv.conditioning * replicate
        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]

        for prompti, seedi in zip(conditioning_local, seed_increment):
            if prompti == 'SKIP':
                continue
            draw_filename = prompti.strip().replace(' ', '-')
            if fix_seed:
                np.random.seed(cfg.env.rnd_seed + seedi)
                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
            else:
                suffixi = suffix

            if with_ema:
                with netm.ema_scope():
                    x, _ = self.sample(netm, sampler, prompti, **cfgv.sample)
            else:
                x, _ = self.sample(netm, sampler, prompti, **cfgv.sample)

            demo_image = latent2im(netm, x)
            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)

        if eval_cnt is not None:
            print_log('Demo printed for {}'.format(eval_cnt))
        return {}

##################
# eval variation #
##################

class eval_stage_variation(eval_stage):
    @torch.no_grad()
    def sample(self, net, sampler, visual_hint, output_dim, scale, n_samples, ddim_steps, ddim_eta):
        h, w = output_dim
        vh = tvtrans.ToTensor()(PIL.Image.open(visual_hint))[None].to(net.device)
        c = net.get_learned_conditioning(vh)
        c = c.repeat(n_samples, 1, 1)
        uc = None
        if scale != 1.0:
            dummy = torch.zeros_like(vh)
            uc = net.get_learned_conditioning(dummy)
            uc = uc.repeat(n_samples, 1, 1)

        shape = [4, h//8, w//8]
        rv = sampler.sample(
            S=ddim_steps,
            conditioning=c,
            batch_size=n_samples,
            shape=shape,
            verbose=False,
            unconditional_guidance_scale=scale,
            unconditional_conditioning=uc,
            eta=ddim_eta)
        return rv

    def __call__(self, **paras):
        cfg = cfguh().cfg
        cfgv = cfg.eval

        net = paras['net']
        eval_cnt = paras.get('eval_cnt', None)
        fix_seed = cfgv.get('fix_seed', False)

        LRANK = sync.get_rank('local')
        LWSIZE = sync.get_world_size('local')

        image_path = self.get_image_path()
        self.create_dir(image_path)
        eval_cnt = paras.get('eval_cnt', None)
        suffix='' if eval_cnt is None else '_'+str(eval_cnt)

        if isinstance(net, (torch.nn.DataParallel,
                            torch.nn.parallel.DistributedDataParallel)):
            netm = net.module
        else:
            netm = net

        with_ema = getattr(netm, 'model_ema', None) is not None
        sampler = self.sampler(netm)
        setattr(netm, 'device', LRANK) # Trick

        color_adj = cfguh().cfg.eval.get('color_adj', False)
        color_adj_keep_ratio = cfguh().cfg.eval.get('color_adj_keep_ratio', 0.5)
        color_adj_simple = cfguh().cfg.eval.get('color_adj_simple', True)

        replicate = cfgv.get('replicate', 1)
        conditioning = cfgv.conditioning * replicate
        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]

        for ci, seedi in zip(conditioning_local, seed_increment):
            if ci == 'SKIP':
                continue

            draw_filename = osp.splitext(osp.basename(ci))[0]

            if fix_seed:
                np.random.seed(cfg.env.rnd_seed + seedi)
                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
            else:
                suffixi = suffix

            if with_ema:
                with netm.ema_scope():
                    x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
            else:
                x, _ = self.sample(netm, sampler, ci, **cfgv.sample)

            demo_image = latent2im(netm, x)
            if color_adj:
                x_adj = []
                for demoi in demo_image:
                    color_adj_f = color_adjust(ref_from=demoi, ref_to=ci)
                    xi_adj = color_adj_f(demoi, keep=color_adj_keep_ratio, simple=color_adj_simple)
                    x_adj.append(xi_adj)
                demo_image = x_adj
            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)

        if eval_cnt is not None:
            print_log('Demo printed for {}'.format(eval_cnt))
        return {}