import os
import numpy as np
import torch
import torch.nn.functional as F
from torchvision.utils import save_image
from dataloader import get_loader
from models.model_main import ModelMain
from models.transformers import denumericalize
from options import get_parser_main_model
from data_utils.svg_utils import render
from models.util_funcs import svg2img, cal_iou
from tqdm import tqdm
from PIL import Image

def test_main_model(opts):
    if opts.streamlit:
        import streamlit as st

    if opts.dir_res:
        os.mkdir(os.path.join(opts.dir_res, "results"))
        dir_res = os.path.join(opts.dir_res, "results")
    else:
        dir_res = os.path.join(f"{opts.exp_path}", "experiments/", opts.name_exp, "results")

    test_loader = get_loader(opts.data_root, opts.img_size, opts.language, opts.char_num, opts.max_seq_len, opts.dim_seq, opts.batch_size, 'test')
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    if opts.streamlit:
        st.write("Loading Model Weight...")
    model_main = ModelMain(opts)
    path_ckpt = os.path.join(f"{opts.model_path}")
    model_main.load_state_dict(torch.load(path_ckpt)['model'])
    model_main.to(device)
    model_main.eval()
    with torch.no_grad():
        
        for test_idx, test_data in enumerate(test_loader):
            for key in test_data: test_data[key] = test_data[key].to(device)

            print("testing font %04d ..."%test_idx)

            dir_save = os.path.join(dir_res, "%04d"%test_idx)
            if not os.path.exists(dir_save):
                os.mkdir(dir_save)
                os.mkdir(os.path.join(dir_save, "imgs"))
                os.mkdir(os.path.join(dir_save, "svgs_single"))
                os.mkdir(os.path.join(dir_save, "svgs_merge"))
            svg_merge_dir = os.path.join(dir_save, "svgs_merge")

            iou_max = np.zeros(opts.char_num)
            idx_best_sample = np.zeros(opts.char_num)

            # syn_svg_merge_f = open(os.path.join(svg_merge_dir, f"{opts.name_ckpt}_syn_merge_{test_idx}_rand_{sample_idx}.html"), 'w') 
            syn_svg_merge_f = open(os.path.join(svg_merge_dir, f"{opts.name_ckpt}_syn_merge_{test_idx}.html"), 'w') 
    
            for sample_idx in tqdm(range(opts.n_samples)):
                
                ret_dict_test, loss_dict_test = model_main(test_data, mode='test')

                svg_sampled = ret_dict_test['svg']['sampled_1']
                sampled_svg_2 = ret_dict_test['svg']['sampled_2']

                img_trg = ret_dict_test['img']['trg']
                img_output = ret_dict_test['img']['out']
                trg_seq_gt = ret_dict_test['svg']['trg']

                img_sample_merge = torch.cat((img_trg.data, img_output.data), -2)
                save_file_merge = os.path.join(dir_save, "imgs", f"merge_{opts.img_size}.png")
                save_image(img_sample_merge, save_file_merge, nrow=8, normalize=True)    
                if opts.streamlit:
                    st.progress((sample_idx+1)/opts.n_samples, f"Generating Font Sample {sample_idx+1} Please wait...")
                    im = Image.open(save_file_merge)
                    st.image(im, caption='img_sample_merge')

                for char_idx in range(opts.char_num):
                    img_gt = (1.0 - img_trg[char_idx,...]).data
                    save_file_gt = os.path.join(dir_save,"imgs", f"{char_idx:02d}_gt.png")
                    save_image(img_gt, save_file_gt, normalize=True)

                    img_sample = (1.0 - img_output[char_idx,...]).data
                    save_file = os.path.join(dir_save,"imgs", f"{char_idx:02d}_{opts.img_size}.png")
                    save_image(img_sample, save_file, normalize=True)

                # write results w/o parallel refinement
                svg_dec_out = svg_sampled.clone().detach()
                for i, one_seq in enumerate(svg_dec_out):
                    syn_svg_outfile = os.path.join(os.path.join(dir_save, "svgs_single"), f"syn_{i:02d}_{sample_idx}_wo_refine.svg")

                    syn_svg_f_ = open(syn_svg_outfile, 'w')
                    try:
                        svg = render(one_seq.cpu().numpy())
                        syn_svg_f_.write(svg)
                        # syn_svg_merge_f.write(svg)
                        if i > 0 and i % 13 == 12:
                            syn_svg_f_.write('<br>')
                            # syn_svg_merge_f.write('<br>')
                        
                    except:
                        continue
                    syn_svg_f_.close()
                
                # write results w/ parallel refinement
                svg_dec_out = sampled_svg_2.clone().detach()
                for i, one_seq in enumerate(svg_dec_out):
                    syn_svg_outfile = os.path.join(os.path.join(dir_save, "svgs_single"), f"syn_{i:02d}_{sample_idx}_refined.svg")

                    syn_svg_f = open(syn_svg_outfile, 'w')
                    try:
                        svg = render(one_seq.cpu().numpy())
                        syn_svg_f.write(svg)
                        #syn_svg_merge_f.write(svg)
                        
                        #if i > 0 and i % 13 == 12:
                        #    syn_svg_merge_f.write('<br>')
                    except:
                        continue
                    syn_svg_f.close()
                    syn_img_outfile = syn_svg_outfile.replace('.svg', '.png')
                    svg2img(syn_svg_outfile, syn_img_outfile, img_size=opts.img_size)
                    iou_tmp, l1_tmp = cal_iou(syn_img_outfile, os.path.join(dir_save, "imgs", f"{i:02d}_{opts.img_size}.png"))
                    iou_tmp = iou_tmp
                    if iou_tmp > iou_max[i]:
                        iou_max[i] = iou_tmp
                        idx_best_sample[i] = sample_idx

            for i in range(opts.char_num):
                # print(idx_best_sample[i])
                syn_svg_outfile_best = os.path.join(os.path.join(dir_save, "svgs_single"), f"syn_{i:02d}_{int(idx_best_sample[i])}_refined.svg")
                syn_svg_merge_f.write(open(syn_svg_outfile_best, 'r').read())
                if i > 0 and i % 13 == 12:
                    syn_svg_merge_f.write('<br>')

            svg_target = trg_seq_gt.clone().detach()
            tgt_commands_onehot = F.one_hot(svg_target[:, :, :1].long(), 4).squeeze()
            tgt_args_denum = denumericalize(svg_target[:, :, 1:])
            svg_target = torch.cat([tgt_commands_onehot, tgt_args_denum], dim=-1)

            for i, one_gt_seq in enumerate(svg_target):
                # gt_svg_outfile = os.path.join(os.path.join(dir_save, "svgs_single"), f"gt_{i:02d}.svg")
                # gt_svg_f = open(gt_svg_outfile, 'w')
                gt_svg = render(one_gt_seq.cpu().numpy())
                # gt_svg_f.write(gt_svg)
                syn_svg_merge_f.write(gt_svg)
                # gt_svg_f.close()
                if i > 0 and i % 13 == 12:
                    syn_svg_merge_f.write('<br>')

            syn_svg_merge_f.close()

        return im

def main():
    
    opts = get_parser_main_model().parse_args()
    opts.name_exp = opts.name_exp + '_' + opts.model_name
    experiment_dir = os.path.join(f"{opts.exp_path}","experiments", opts.name_exp)
    print(f"Testing on experiment {opts.name_exp}...")
    # Dump options
    test_main_model(opts)

if __name__ == "__main__":
    main()