#!/usr/bin/env python

import numpy as np
import pandas as pd

from sklearn.svm import SVC
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.model_selection import train_test_split

from tqdm import tqdm
import random
from os.path import join
import os
import pickle

import torch

import matplotlib.pyplot as plt
import PIL
from PIL import Image, ImageColor

import sys
sys.path.append('backend')
from color_annotations import extract_color
from networks_stylegan3 import *
sys.path.append('.')

import dnnlib 
import legacy

def hex2rgb(hex_value):
    h = hex_value.strip("#") 
    rgb = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
    return rgb

def rgb2hsv(r, g, b):
    # Normalize R, G, B values
    r, g, b = r / 255.0, g / 255.0, b / 255.0
    
    # h, s, v = hue, saturation, value
    max_rgb = max(r, g, b)    
    min_rgb = min(r, g, b)   
    difference = max_rgb-min_rgb 
    
    # if max_rgb and max_rgb are equal then h = 0
    if max_rgb == min_rgb:
        h = 0
    
    # if max_rgb==r then h is computed as follows
    elif max_rgb == r:
        h = (60 * ((g - b) / difference) + 360) % 360
    
    # if max_rgb==g then compute h as follows
    elif max_rgb == g:
        h = (60 * ((b - r) / difference) + 120) % 360
    
    # if max_rgb=b then compute h
    elif max_rgb == b:
        h = (60 * ((r - g) / difference) + 240) % 360
    
    # if max_rgb==zero then s=0
    if max_rgb == 0:
        s = 0
    else:
        s = (difference / max_rgb) * 100
    
    # compute v
    v = max_rgb * 100
    # return rounded values of H, S and V
    return tuple(map(round, (h, s, v)))
 

class DisentanglementBase:
    def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s=False, variable='H1', categorical=True):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print('Using device', self.device)
        self.repo_folder = repo_folder
        self.model = model.to(self.device)
        self.annotations = annotations
        self.df = df
        self.space = space
        self.categorical = categorical
        self.variable = variable
        
        self.layers = ['input', 'L0_36_512', 'L1_36_512', 'L2_36_512', 'L3_52_512',
                       'L4_52_512', 'L5_84_512', 'L6_84_512', 'L7_148_512', 'L8_148_512', 
                       'L9_148_362', 'L10_276_256', 'L11_276_181', 'L12_276_128', 
                       'L13_256_128', 'L14_256_3']
        self.layers_shapes = [4, 512, 512, 512, 512, 512, 512, 512, 512, 512, 512, 362, 256, 181, 128, 128]
        self.decoding_layers = 16
        self.colors_list = colors_list
        
        self.to_hsv()
        if compute_s:
            self.get_s_space()
        

    def to_hsv(self):
        """
        The tohsv function takes the top 3 colors of each image and converts them to HSV values.
        It then adds these values as new columns in the dataframe.
        
        :param self: Allow the function to access the dataframe
        :return: The dataframe with the new columns added
        :doc-author: Trelent
        """
        print('Adding HSV encoding')
        self.df['H1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
        self.df['H2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
        self.df['H3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
        
        self.df['S1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
        self.df['S2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
        self.df['S3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
        
        self.df['V1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
        self.df['V2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
        self.df['V3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
        
        print('Adding RGB encoding')
        self.df['R1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
        self.df['R2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
        self.df['R3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
        
        self.df['G1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
        self.df['G2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
        self.df['G3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
        
        self.df['B1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
        self.df['B2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
        self.df['B3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
        return self.df
    
    def get_encoded_latent(self):
        # ... (existing code for getX)
        if self.space.lower() == 'w':
            X = np.array(self.annotations['w_vectors']).reshape((len(self.annotations['w_vectors']), 512))
        elif self.space.lower() == 'z':
            X = np.array(self.annotations['z_vectors']).reshape((len(self.annotations['z_vectors']), 512))
        elif self.space.lower() == 's':
            concat_v = []
            for i in range(len(self.annotations['w_vectors'])):
                concat_v.append(np.concatenate(self.annotations['s_vectors'][i], axis=1))
            X = np.array(concat_v)
            X = X[:, 0, :]
        else:
            Exception("Sorry, option not available, select among Z, W, S")
            
        print('Shape embedding:', X.shape)
        return X
    
    def get_train_val(self, extremes=False):
        X = self.get_encoded_latent()
        y = np.array(self.df[self.variable].values)
        if self.categorical:
            y_cat = pd.cut(y, 
                            bins=[x * 360 / len(self.colors_list) if x < len(self.colors_list) 
                                  else 360 for x in range(len(self.colors_list) + 1)],
                            labels=self.colors_list
                            ).fillna(self.colors_list[0])
            x_train, x_val, y_train, y_val = train_test_split(X, y_cat, test_size=0.2)
        else:
            if extremes:
                # Calculate the number of elements to consider (10% of array size)
                num_elements = int(0.2 * len(y))
                # Get indices of the top num_elements maximum values
                top_indices = np.argpartition(array, -num_elements)[-num_elements:]
                bottom_indices = np.argpartition(array, -num_elements)[:num_elements]
                y_ext = y[top_indices + bottom_indices, :]
                X_ext = X[top_indices + bottom_indices, :]
                x_train, x_val, y_train, y_val = train_test_split(X_ext, y_ext, test_size=0.2)
            else:
                x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.2)
        return x_train, x_val, y_train, y_val
    
    def generate_orig_image(self, vec, seed=False):
        """
        The generate_original_image function takes in a latent vector and the model,
        and returns an image generated from that latent vector.
        
        
        :param z: Generate the image
        :param model: Generate the image
        :return: A pil image
        :doc-author: Trelent
        """
        G = self.model.to(self.device) # type: ignore
        # Labels.
        label = torch.zeros([1, G.c_dim], device=self.device)
        if seed:
            seed = vec
            vec = self.annotations['z_vectors'][seed]
            
        Z = torch.from_numpy(vec.copy()).to(self.device)
        img = G(Z, label, truncation_psi=1, noise_mode='const')
        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
        img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
        return img
            
def main():
    repo_folder = '.'
    annotations_file = join(repo_folder, 'data/textile_annotated_files/seeds0000-100000_S.pkl')
    with open(annotations_file, 'rb') as f:
        annotations = pickle.load(f)

    df_file = join(repo_folder, 'data/textile_annotated_files/top_three_colours.csv')
    df = pd.read_csv(df_file).fillna('#000000')

    model_file = join(repo_folder, 'data/textile_model_files/network-snapshot-005000.pkl')
    with dnnlib.util.open_url(model_file) as f:
        model = legacy.load_network_pkl(f)['G_ema'] # type: ignore

    colors_list = ['Red', 'Orange', 'Yellow', 'Yellow Green', 'Chartreuse Green',
                   'Kelly Green', 'Green Blue Seafoam', 'Cyan Blue',
                   'Warm Blue', 'Indigo', 'Purple Magenta', 'Magenta Pink']
    colors_list = ['Red Orange', 'Yellow', 'Green', 'Light Blue',
                   'Blue', 'Purple', 'Pink']
    
    
    disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space='w', colors_list=colors_list)
    # x_train, x_val, y_train, y_val = disentanglemnet_exp.get_train_val()
    # print(colors_list)
    # print(np.unique(y_train, return_counts=True))
    
    
    # for i, color in enumerate(colors_list):
    #     idxs = np.where(y_train == color)
    #     x_color = x_train[idxs][:30, :]
    #     print(x_color.shape)
    #     print('Generating images of color ' + color)
    #     for j, vec in enumerate(x_color):
    #         vec = np.expand_dims(vec, axis=0)
    #         img = disentanglemnet_exp.generate_orig_image(vec)
    #         img.save(f'{repo_folder}/colors_test/color_{color}_{j}.png')
                                            
    df = disentanglemnet_exp.to_hsv()
    df['color'] = pd.cut(df['H1'], 
                        bins=[x * 360 / len(colors_list) if x < len(colors_list) 
                              else 360 for x in range(len(colors_list) + 1)],
                            labels=colors_list
                            ).fillna(colors_list[0])  
    
    print(df['color'].value_counts())
    df['seed'] = df['fname'].str.split('/').apply(lambda x: x[-1]).str.replace('seed', '').str.replace('.png','').astype(int)
    print(df[df['seed'] == 3][['H1', 'S1', 'V1', 'R1', 'B1', 'G1']])
    for i, color in enumerate(colors_list):
        idxs = df['color'] == color
        x_color = df['seed'][idxs][:30]
        print('Generating images of color ' + color)
        for j, vec in enumerate(x_color):
            img = disentanglemnet_exp.generate_orig_image(int(vec), seed=True)
            img.save(f'{repo_folder}/colors_test/color_{color}_{j}corrected.png')
            
if __name__ == "__main__":
    main()