import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from torchvision.utils import save_image
from tensorboardX import SummaryWriter
from dataloader import get_loader
from models.image_encoder import ImageEncoder
from models.image_decoder import ImageDecoder
from models.modality_fusion import ModalityFusion
from models.vgg_perceptual_loss import VGGPerceptualLoss
from models.svg_decoder import SVGLSTMDecoder, SVGMDNTop
from models.svg_encoder import SVGLSTMEncoder
from models import util_funcs
from options import get_parser_main_model
from data_utils.svg_utils import render
from models.imgsr.modules import TrainOptions, create_model

opts = get_parser_main_model().parse_args()
opts.experiment_name = opts.experiment_name + '_' + opts.model_name
os.makedirs("experiments", exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def get_z(temperature=1.0):
    z = torch.randn(1, 128).to(device) * temperature
    return z

def get_img_feat(temperature=0.5):
    z = torch.randn(1, 1024).to(device) * temperature
    return z
def get_enc_z(temperature=0.5):
    z = torch.randn(4, 1024).to(device) * temperature
    return z


svg_dir = "output"

exp_dir = os.path.join("experiments", opts.experiment_name)
ckpt_dir = os.path.join(exp_dir, "checkpoints")
res_dir = os.path.join(exp_dir, "results")

svg_decoder = SVGLSTMDecoder(char_categories=opts.char_categories,
                             bottleneck_bits=opts.bottleneck_bits, mode=opts.mode, max_sequence_length=opts.max_seq_len,
                             hidden_size=opts.hidden_size,
                             num_hidden_layers=opts.num_hidden_layers,
                             feature_dim=opts.seq_feature_dim, ff_dropout=opts.ff_dropout, rec_dropout=opts.rec_dropout)

mdn_top_layer = SVGMDNTop(num_mixture=opts.num_mixture, seq_len=opts.max_seq_len, hidden_size=opts.hidden_size,
                          mode=opts.mode, mix_temperature=opts.mix_temperature,
                          gauss_temperature=opts.gauss_temperature, dont_reduce=opts.dont_reduce_loss)

modality_fusion = ModalityFusion(img_feat_dim = 16 * opts.image_size, hidden_size = opts.hidden_size, ref_nshot = opts.ref_nshot, bottleneck_bits = opts.bottleneck_bits, mode=opts.mode)


# load parameters
epoch = opts.test_epoch

svg_decoder_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.seqdec.pth")
svg_decoder.load_state_dict(torch.load(svg_decoder_fpath, map_location=device))
svg_decoder.eval()


mdn_top_layer_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.mdntl.pth")
mdn_top_layer.load_state_dict(torch.load(mdn_top_layer_fpath, map_location=device))
mdn_top_layer.eval()


modality_fusion_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.modalfuse.pth")
modality_fusion.load_state_dict(torch.load(modality_fusion_fpath, map_location=device))
modality_fusion.eval()

# to device
modality_fusion = modality_fusion.to(device)
svg_decoder = svg_decoder.to(device)
mdn_top_layer = mdn_top_layer.to(device)

val_img_l1_loss = 0.0
val_img_pt_loss = 0.0
mean = np.load('./data/mean.npz')
std = np.load('./data/stdev.npz')

mean = torch.from_numpy(mean).to(device).to(torch.float32)
std = torch.from_numpy(std).to(device).to(torch.float32)

latent_feat = get_z()
latent_feat = latent_feat.repeat(opts.char_categories, 1)

trg_char = F.one_hot(torch.arange(52), num_classes=opts.char_categories)
sd_init_state = svg_decoder.init_state_input(latent_feat, trg_char)
hidden, cell = sd_init_state['hidden'], sd_init_state['cell']
hidden_self, cell_self = hidden, cell
# outputs_self = torch.zeros(trg_seq.size(0), trg_seq.size(1), opts.hidden_size).to(device)


tgt_len = 51
trg_seq_size_1 = 52
sampled_svg = torch.zeros(tgt_len, trg_seq_size_1, opts.seq_feature_dim).to(device)
for t in range(0, tgt_len):
    # self sample results
    if t == 0:
        inpt_self = torch.zeros(trg_seq_size_1, opts.seq_feature_dim).to(device)
    else:
        inpt_self = sampled_svg[t - 1]
    decoder_output_self = svg_decoder(inpt_self, hidden_self, cell_self)
    output_self, hidden_self, cell_self = decoder_output_self['output'], decoder_output_self['hidden'], \
                                          decoder_output_self['cell']
    top_output_self = mdn_top_layer(output_self)
    sampled_step = mdn_top_layer.sample(top_output_self, output_self, opts.mode)
    #print(sampled_step.size())
    sampled_svg[t] = sampled_step

svg_dec_out = sampled_svg.clone().detach()
svg_dec_out = svg_dec_out.transpose(0, 1)
svg_dec_out = svg_dec_out * std + mean

for i, one_seq in enumerate(svg_dec_out):
    syn_svg_outfile = os.path.join(svg_dir, f"generated{i}.svg")
    syn_svg_f = open(syn_svg_outfile, 'w')
    try:
        svg = render(one_seq.cpu().numpy())
        syn_svg_f.write(svg)
    except:
        continue
    syn_svg_f.close()

syn_svg_f.close()