import os
import time
import json
import base64
import argparse
import importlib
from glob import glob
from PIL import Image
from imageio import imsave

import torch
import torchvision.utils as vutils

import sys
sys.path.append(".")

import numpy as np
from libs.test_base import TesterBase
from libs.utils import colorEncode, label2one_hot_torch
from tqdm import tqdm
from libs.options import BaseOptions
import torch.nn.functional as F
from libs.nnutils import poolfeat, upfeat

import streamlit as st
from skimage.segmentation import slic
import torchvision.transforms.functional as TF
import torchvision.transforms as transforms
from st_clickable_images import clickable_images

args = BaseOptions().gather_options()
if args.img_path is not None:
    args.exp_name = os.path.join(args.exp_name, args.img_path.split('/')[-1].split('.')[0])
args.batch_size = 1
args.data_path = "/home/xli/DATA/BSR_processed/train"
args.label_path = "/home/xli/DATA/BSR/BSDS500/data/groundTruth"
args.device = torch.device("cpu")
args.nsamples = 500
args.out_dir = os.path.join('cachedir', args.exp_name)
os.makedirs(args.out_dir, exist_ok=True)
args.global_code_ch = args.hidden_dim
args.netG_use_noise = True
args.test_time = (args.test_time == 1)

if not hasattr(args, 'tex_code_dim'):
    args.tex_code_dim = 256

class Tester(TesterBase):
    def define_model(self):
        """Define model
        """
        args = self.args
        module = importlib.import_module('models.week0417.{}'.format(args.model_name))
        self.model = module.AE(args)
        self.model.to(args.device)
        self.model.eval()
        return

    def draw_color_seg(self, seg):
        seg = seg.detach().cpu().numpy()
        color_ = []
        for i in range(seg.shape[0]):
            colori = colorEncode(seg[i].squeeze())
            colori = torch.from_numpy(colori / 255.0).float().permute(2, 0, 1)
            color_.append(colori)
        color_ = torch.stack(color_)
        return color_

    def to_pil(self, tensor):
        return transforms.ToPILImage()(tensor.cpu().squeeze().clamp(0.0, 1.0)).convert("RGB")

    def display_synthesis(self):
        with st.spinner('Running...'):
            with torch.no_grad():
                grouping_mask = self.model_forward_synthesis(self.data, self.slic, return_type = 'grouping')

            data = (self.data + 1) / 2.0

            seg = grouping_mask.view(-1, 1, args.crop_size, args.crop_size)
            color_vq = self.draw_color_seg(seg)
            color_vq = color_vq * 0.8 + data.cpu() * 0.2

            st.markdown('<p class="big-font">Given the image you chose, our model decomposes the image into ten texture segments, each depicts one kind of texture in the image.</p>', unsafe_allow_html=True)
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.markdown("")

            with col2:
                st.markdown("Chosen image")
                st.image(self.to_pil(data))

            with col3:
                st.markdown("Grouping mask")
                st.image(self.to_pil(color_vq))

            with col4:
                st.markdown("")

        seg_onehot = label2one_hot_torch(seg, C = 10)
        parts = data.cpu() * seg_onehot.squeeze().unsqueeze(1)

        st.markdown('<p class="big-font">We show all texture segments below. To synthesize an arbitrary-sized texture image from a texture segment, choose and click one of the texture segments below.</p>', unsafe_allow_html=True)
        tmp_img_list = []
        for i in range(parts.shape[0]):
            part_img = self.to_pil(parts[i])
            out_path = 'tmp/{}.png'.format(i)
            part_img.save(out_path)

            with open(out_path, "rb") as image:
                encoded = base64.b64encode(image.read()).decode()
                tmp_img_list.append(f"data:image/jpeg;base64,{encoded}")

        tex_idx = clickable_images(
            tmp_img_list,
            titles=[f"Group #{str(i)}" for i in range(len(tmp_img_list))],
            div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
            img_style={"margin": "5px", "height": "150px"},
            key=0
        )

        if tex_idx > -1:
            with st.spinner('Running...'):
                st.markdown('<p class="big-font">You can slide the bar below to set the size of the synthesized texture image.</p>', unsafe_allow_html=True)
                tex_size = st.slider('', 0, 1000, 256)
                tex_size = (tex_size // 8) * 8
                with torch.no_grad():
                    tex = self.model_forward_synthesis(self.data, self.slic, tex_idx = tex_idx, tex_size = tex_size, return_type = 'tex')
                    col1, col2, col3, col4 = st.columns([1, 1, 4, 1])
                    with col1:
                        st.markdown("")

                    with col2:
                        st.markdown("Chosen examplar segment")
                        st.image(self.to_pil(parts[tex_idx]))

                    with col3:
                        st.markdown("Synthesized texture image")
                        st.image(self.to_pil(tex))

                    with col4:
                        st.markdown("")
            st.markdown('<p class="big-font">You can choose another image from the examplar images on the top and start again!</p>', unsafe_allow_html=True)

    def model_forward_synthesis(self, rgb_img, slic, epoch = 1000, test_time = False,
                                test = True, tex_idx = None, tex_size = 256,
                                return_type = 'tex', fill_idx = None, remove_idx = None):
        args = self.args
        B, _, imgH, imgW = rgb_img.shape

        # Encoder: img (B, 3, H, W) -> feature (B, C, imgH//8, imgW//8)
        conv_feat, _ = self.model.enc(rgb_img)
        B, C, H, W = conv_feat.shape

        # Texture code for each superpixel
        tex_code = self.model.ToTexCode(conv_feat)

        code = F.interpolate(tex_code, size = (imgH, imgW), mode = 'bilinear', align_corners = False)
        pool_code = poolfeat(code, slic, avg = True)

        prop_code, sp_assign, conv_feats = self.model.gcn(pool_code, slic, (args.add_clustering_epoch <= epoch))
        softmax = F.softmax(sp_assign * args.temperature, dim = 1)
        if return_type == 'grouping':
            return torch.argmax(sp_assign.cpu(), dim = 1)

        tex_seg = poolfeat(conv_feats, softmax, avg = True)
        seg = label2one_hot_torch(torch.argmax(softmax, dim = 1).unsqueeze(1), C = softmax.shape[1])

        sampled_code = tex_seg[:, tex_idx, :]
        rec_tex = sampled_code.view(1, -1, 1, 1).repeat(1, 1, tex_size, tex_size)
        sine_wave = self.model.get_sine_wave(rec_tex, 'rec')[:1]
        H = tex_size // 8; W = tex_size // 8
        noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
        dec_input = torch.cat((sine_wave, noise), dim = 1)

        #weight = self.model.ChannelWeight(rec_tex)
        #weight = F.adaptive_avg_pool2d(weight, output_size = (1)).view(weight.shape[0], -1, 1, 1)
        #weight = torch.sigmoid(weight)
        #dec_input *= weight

        rep_rec = self.model.G(dec_input, rec_tex)
        rep_rec = (rep_rec + 1) / 2.0
        return rep_rec

    def display_editing(self):
        with st.spinner('Running...'):
            with torch.no_grad():
                grouping_mask = self.model_forward_editing(self.data, self.slic, return_type = 'grouping')

            data = (self.data + 1) / 2.0

            seg = grouping_mask.view(-1, 1, args.crop_size, args.crop_size)
            color_vq = self.draw_color_seg(seg)
            color_vq = color_vq * 0.8 + data.cpu() * 0.2

            st.markdown('<p class="big-font">Given the image you chose, our model decomposes the image into ten texture segments, each depicts one kind of texture in the image.</p>', unsafe_allow_html=True)
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.markdown("")

            with col2:
                st.markdown("Chosen image")
                st.image(self.to_pil(data))

            with col3:
                st.markdown("Grouping mask")
                st.image(self.to_pil(color_vq))

            with col4:
                st.markdown("")

        seg_onehot = label2one_hot_torch(seg, C = 10)
        parts = data.cpu() * seg_onehot.squeeze().unsqueeze(1)

        st.markdown('<p class="big-font">We show all texture segments below.</p>', unsafe_allow_html=True)
        tmp_img_list = []
        for i in range(parts.shape[0]):
            part_img = self.to_pil(parts[i])
            out_path = 'tmp/{}.png'.format(i)
            part_img.save(out_path)

            with open(out_path, "rb") as image:
                encoded = base64.b64encode(image.read()).decode()
                tmp_img_list.append(f"data:image/jpeg;base64,{encoded}")

        tex_idx = clickable_images(
            tmp_img_list,
            titles=[f"Group #{str(i)}" for i in range(len(tmp_img_list))],
            div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
            img_style={"margin": "5px", "height": "150px"},
            key=2
        )

        st.markdown('<p class="big-font">Choose one mask for texture editing.</p>', unsafe_allow_html=True)
        mask_list = glob(os.path.join("data/masks/*.png"))
        byte_mask_list = []
        for img_path in mask_list:
            seg = Image.open(img_path).convert("L")
            seg = np.asarray(seg)
            seg = torch.from_numpy(seg).view(1, 1, seg.shape[0], seg.shape[1])
            color_vq = self.draw_color_seg(seg)
            vutils.save_image(color_vq, 'tmp/tmp.png')
            with open('tmp/tmp.png', "rb") as image:
                encoded = base64.b64encode(image.read()).decode()
                byte_mask_list.append(f"data:image/jpeg;base64,{encoded}")
        img_idx = clickable_images(
            byte_mask_list,
            titles=[f"Group #{str(i)}" for i in range(len(byte_mask_list))],
            div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
            img_style={"margin": "5px", "height": "150px"},
        )
        mask_path = mask_list[img_idx]

        st.markdown('<p class="big-font">Choose the texture segment for each group in the given mask below.</p>', unsafe_allow_html=True)
        given_mask = Image.open(mask_path).convert("L")
        given_mask = np.asarray(given_mask)
        given_mask = torch.from_numpy(given_mask)
        H, W = given_mask.shape[0], given_mask.shape[1]
        given_mask = label2one_hot_torch(given_mask.view(1, 1, H, W), C = (given_mask.max()+1))
        given_mask = F.interpolate(given_mask, size = (512, 512), mode = 'bilinear', align_corners = False)
        mask_img_list = []
        for i in range(given_mask.shape[1]):
            part_img = self.to_pil(given_mask[0, i])
            out_path = 'tmp/{}.png'.format(i)
            part_img.save(out_path)

            with open(out_path, "rb") as image:
                encoded = base64.b64encode(image.read()).decode()
                mask_img_list.append(f"data:image/jpeg;base64,{encoded}")

        part_idx = clickable_images(
            mask_img_list,
            div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
            img_style={"margin": "5px", "height": "150px"},
            key=1
        )

        cols = st.columns(len(mask_img_list))
        options = []
        for i, col in enumerate(cols):
            with col:
                option = st.selectbox(
                        "",
                        ([str(ii) for ii in range(10)]),
                        key = i)
                options.append(int(option))
        print(options)

        if len(options) > 0:
            with st.spinner('Running...'):
                st.markdown('<p class="big-font">Edited image is shown below.</p>', unsafe_allow_html=True)
                #tex_size = st.slider('', 0, 1000, 256)
                #tex_size = (tex_size // 8) * 8
                with torch.no_grad():
                    edited = self.model_forward_editing(self.data, self.slic, options=options, given_mask=given_mask, return_type = 'edited')
                    col1, col2, col3 = st.columns([1, 1, 1])
                    with col1:
                        st.markdown("Input image")
                        img = F.interpolate(self.data, size = edited.shape[-2:], mode = 'bilinear', align_corners = False)
                        st.image(self.to_pil((img + 1) / 2.0))
                        print(img.shape, edited.shape)

                    with col2:
                        st.markdown("Given mask")
                        seg = Image.open(mask_path).convert("L")
                        seg = np.asarray(seg)
                        seg = torch.from_numpy(seg).view(1, 1, seg.shape[0], seg.shape[1])
                        color_vq = self.draw_color_seg(seg)
                        color_vq = F.interpolate(color_vq, size = (512, 512), mode = 'bilinear', align_corners = False)
                        st.image(self.to_pil(color_vq))

                    with col3:
                        st.markdown("Synthesized texture image")
                        st.image(self.to_pil(edited))

            st.markdown('<p class="big-font">You can choose another image from the examplar images on the top and start again!</p>', unsafe_allow_html=True)

    def model_forward_editing(self, rgb_img, slic, epoch = 1000, test_time = False,
                      test = True, tex_idx = None, tex_size = 256,
                      return_type = 'edited', fill_idx = None, remove_idx = None,
                      options = None, given_mask = None):
        args = self.args
        B, _, imgH, imgW = rgb_img.shape

        # Encoder: img (B, 3, H, W) -> feature (B, C, imgH//8, imgW//8)
        conv_feat, _ = self.model.enc(rgb_img)
        B, C, H, W = conv_feat.shape

        # Texture code for each superpixel
        tex_code = self.model.ToTexCode(conv_feat)

        code = F.interpolate(tex_code, size = (imgH, imgW), mode = 'bilinear', align_corners = False)
        pool_code = poolfeat(code, slic, avg = True)

        prop_code, sp_assign, conv_feats = self.model.gcn(pool_code, slic, (args.add_clustering_epoch <= epoch))
        softmax = F.softmax(sp_assign * args.temperature, dim = 1)
        if return_type == 'grouping':
            return torch.argmax(sp_assign.cpu(), dim = 1)

        tex_seg = poolfeat(conv_feats, softmax, avg = True)
        seg = label2one_hot_torch(torch.argmax(softmax, dim = 1).unsqueeze(1), C = softmax.shape[1])

        rec_tex = torch.zeros((1, tex_seg.shape[-1], 512, 512))
        for i in range(given_mask.shape[1]):
            label = options[i]
            code = tex_seg[0, label, :].view(1, -1, 1, 1).repeat(1, 1, 512, 512)
            rec_tex += code * given_mask[:, i:i+1]
        tex_size = 512
        sine_wave = self.model.get_sine_wave(rec_tex, 'rec')[:1]
        H = tex_size // 8; W = tex_size // 8
        noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
        dec_input = torch.cat((sine_wave, noise), dim = 1)

        rep_rec = self.model.G(dec_input, rec_tex)
        rep_rec = (rep_rec + 1) / 2.0
        return rep_rec

    def load_data(self, data_path):
        rgb_img = Image.open(data_path)
        crop_size = self.args.crop_size
        i = 40; j = 40; h = crop_size; w = crop_size
        rgb_img = transforms.Resize(size=320)(rgb_img)
        rgb_img = TF.crop(rgb_img, i, j, h, w)

        # compute superpixel
        sp_num = 196
        slic_i = slic(np.array(rgb_img), n_segments=sp_num, compactness=10, start_label=0, min_size_factor=0.3)
        slic_i = torch.from_numpy(slic_i)
        slic_i[slic_i >= sp_num] = sp_num - 1
        oh = label2one_hot_torch(slic_i.unsqueeze(0).unsqueeze(0), C = sp_num).squeeze()
        self.slic = oh.unsqueeze(0).to(args.device)

        rgb_img = TF.to_tensor(rgb_img)
        rgb_img = rgb_img.unsqueeze(0)
        self.data = rgb_img.to(args.device) * 2 - 1

    def load_model(self, model_path):
        self.model = torch.nn.DataParallel(self.model)
        cpk = torch.load(model_path, map_location=torch.device('cpu'))
        saved_state_dict = cpk['model']
        self.model.load_state_dict(saved_state_dict)
        self.model = self.model.module
        return

    """
    def test(self):
        #for iteration in tqdm(range(args.nsamples)):
        self.test_step(0)
        self.display(0, 'train')
    """

def main():
    #torch.cuda.empty_cache()
    st.set_page_config(layout="wide")
    st.markdown("""
                <style>
                .big-font {
                    font-size:30px !important;
                }
                </style>
                """, unsafe_allow_html=True)

    st.title("Scraping Textures from Natural Images for Synthesis and Editing")
    #st.markdown("**In this demo, we show how to scrape textures from natural images for texture synthesis and editing.**")
    st.markdown('<p class="big-font">In this demo, we show how to scrape textures from natural images for texture synthesis and editing.</p>', unsafe_allow_html=True)
    st.markdown("## Texture synthesis")
    st.markdown('<p class="big-font">Here we provide a set of example images, please choose and click one image to start.</p>', unsafe_allow_html=True)
    img_list = glob(os.path.join("data/images/*.jpg"))
    test_img_list = glob(os.path.join("data/test_images/*.jpg"))
    img_list.extend(test_img_list)
    byte_img_list = []
    for img_path in img_list:
        with open(img_path, "rb") as image:
            encoded = base64.b64encode(image.read()).decode()
            byte_img_list.append(f"data:image/jpeg;base64,{encoded}")
    img_idx = clickable_images(
        byte_img_list,
        titles=[f"Group #{str(i)}" for i in range(len(byte_img_list))],
        div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
        img_style={"margin": "5px", "height": "150px"},
    )
    img_path = img_list[img_idx]

    img_name = img_path.split("/")[-1]
    args.pretrained_path = os.path.join("weights/{}/cpk.pth".format(img_name.split(".")[0]))

    if img_idx > -1:
        tester = Tester(args)
        tester.define_model()
        tester.load_data(img_path)
        tester.load_model(args.pretrained_path)
        app_idx = st.selectbox('Please select between texture synthesis or editing',
                              ["Texture Synthesis", "Texture Editing"])
        if app_idx == 'Texture Editing':
            st.header("Texture Editing")
            tester.display_editing()
        else:
            st.header("Texture Synthesis")
            tester.display_synthesis()

if __name__ == '__main__':
    os.system("pip install torch-geometric==1.7.2")
    main()