import os
import json
import math
import copy
import numpy as np
import networkx as nx
from collections import namedtuple
from minecraft_pb2 import *
from utils_blocks_categories import *
from PIL import ImageColor
from llist import dllist, dllistnode
from typing import Callable, Tuple
from abc import ABC, abstractmethod
from bisect import bisect_left
from colormath.color_objects import *
from colormath import color_diff_matrix
from colormath.color_conversions import convert_color

# region colormath.color_diff 
# implementation form:
#   https://python-colormath.readthedocs.io/en/latest/_modules/colormath/color_diff.html#delta_e_cmc
# depricated np.asscalar(var) replaced by var.item()

# -*- coding: utf-8 -*-
"""
The functions in this module are used for comparing two LabColor objects
using various Delta E formulas.
"""

def _get_lab_color1_vector(color):
    """
    Converts an LabColor into a np vector.

    :param LabColor color:
    :rtype: np.ndarray
    """
    if not color.__class__.__name__ == 'LabColor':
        raise ValueError(
            "Delta E functions can only be used with two LabColor objects.")
    return np.array([color.lab_l, color.lab_a, color.lab_b])


def _get_lab_color2_matrix(color):
    """
    Converts an LabColor into a np matrix.

    :param LabColor color:
    :rtype: np.ndarray
    """
    if not color.__class__.__name__ == 'LabColor':
        raise ValueError(
            "Delta E functions can only be used with two LabColor objects.")
    return np.array([(color.lab_l, color.lab_a, color.lab_b)])


# noinspection PyPep8Naming
def delta_e_cie1976(color1, color2):
    """
    Calculates the Delta E (CIE1976) of two colors.
    """
    color1_vector = _get_lab_color1_vector(color1)
    color2_matrix = _get_lab_color2_matrix(color2)
    delta_e = color_diff_matrix.delta_e_cie1976(color1_vector, color2_matrix)[0]
    return delta_e.item()



# noinspection PyPep8Naming
def delta_e_cie1994(color1, color2, K_L=1, K_C=1, K_H=1, K_1=0.045, K_2=0.015):
    """
    Calculates the Delta E (CIE1994) of two colors.

    K_l:
      0.045 graphic arts
      0.048 textiles
    K_2:
      0.015 graphic arts
      0.014 textiles
    K_L:
      1 default
      2 textiles
    """
    color1_vector = _get_lab_color1_vector(color1)
    color2_matrix = _get_lab_color2_matrix(color2)
    delta_e = color_diff_matrix.delta_e_cie1994(
        color1_vector, color2_matrix, K_L=K_L, K_C=K_C, K_H=K_H, K_1=K_1, K_2=K_2)[0]
    return delta_e.item()



# noinspection PyPep8Naming
def delta_e_cie2000(color1, color2, Kl=1, Kc=1, Kh=1):
    """
    Calculates the Delta E (CIE2000) of two colors.
    """
    color1_vector = _get_lab_color1_vector(color1)
    color2_matrix = _get_lab_color2_matrix(color2)
    delta_e = color_diff_matrix.delta_e_cie2000(
        color1_vector, color2_matrix, Kl=Kl, Kc=Kc, Kh=Kh)[0]
    return delta_e.item()



# noinspection PyPep8Naming
def delta_e_cmc(color1, color2, pl=2, pc=1):
    """
    Calculates the Delta E (CMC) of two colors.

    CMC values
      Acceptability: pl=2, pc=1
      Perceptability: pl=1, pc=1
    """
    color1_vector = _get_lab_color1_vector(color1)
    color2_matrix = _get_lab_color2_matrix(color2)
    delta_e = color_diff_matrix.delta_e_cmc(
        color1_vector, color2_matrix, pl=pl, pc=pc)[0]
    return delta_e.item()

# endregion

class ColorConverter:
    
    @staticmethod
    def rgb_to_str(rgb_color):
        return '#{0:02X}{1:02X}{2:02X}'.format(*rgb_color)
        
    @staticmethod
    def str_to_rgb(str_color):
        return ImageColor.getrgb(str_color)
    
    @staticmethod
    def rgb_to_obj(rgb_color, color_rep):
        return convert_color(sRGBColor(*rgb_color, is_upscaled=True), color_rep)
    
    @staticmethod
    def obj_to_rgb(color):
        return convert_color(color,sRGBColor).get_upscaled_value_tuple()
    
class ColorMetrics:
    
    _hue_max = 360
    @staticmethod
    def value_indermediate(val1, val2, ratio):
        return val1 + ratio*(val2-val1)

    @staticmethod
    def hue_distance(hue1, hue2):
        return (hue2 - hue1 + ColorMetrics._hue_max / 2) % ColorMetrics._hue_max - ColorMetrics._hue_max / 2

    @staticmethod  
    def get_intermediate_color(color1, color2, ratio):
        if not isinstance(color1, type(color2)):
            raise ValueError("Input colors must be of the same type.")
        intermediate_color = []
        for i,(val1, val2) in enumerate(zip(color1.get_value_tuple(), color2.get_value_tuple())):
            intermediate_val = ColorMetrics.value_indermediate(val1, val2, ratio)
            if i==0 and isinstance(color1, HSVColor):
                offset = ColorMetrics.value_indermediate(0, ColorMetrics.hue_distance(val1, val2), ratio)
                intermediate_val = val1 + offset
                if intermediate_val<0:
                    intermediate_val += ColorMetrics._hue_max
            intermediate_color.append(intermediate_val)
        return type(color1)(*intermediate_color)
    
    @staticmethod
    def rgb_distance(color1, color2):
        rgb1 = np.array(color1.get_upscaled_value_tuple())
        rgb2 = np.array(color2.get_upscaled_value_tuple())
        return np.linalg.norm(rgb1 - rgb2)
      
    @staticmethod
    def lab_distance(color1, color2):
        return delta_e_cie2000(color1, color2)
     
    @staticmethod
    def lab_approx_distance(color1, color2):
        lab1 = np.array(color1.get_value_tuple())
        lab2 = np.array(color2.get_value_tuple())
        return np.linalg.norm(lab1 - lab2)
    
    @staticmethod
    def hsv_distance(color1, color2):
        dh = abs(ColorMetrics.hue_distance(color1.hsv_h, color2.hsv_h)) / 180.0
        ds = color1.hsv_s - color2.hsv_s
        dv = color1.hsv_v - color2.hsv_v
        dc = color1.hsv_s * color2.hsv_s * (1 - np.cos(dh * np.pi))  

        return np.sqrt(dh**2 + ds**2 + dv**2 + dc)

class BlockColors:

    _default_file = os.path.join("py_scripts","block_color_extraction","block_to_color.json")
    PathEdge = namedtuple('PathEdge',['node','dist'])

    def __init__(
        self, 
        block_dict: dict[str, Tuple[int, int, int]],
        color_rep = sRGBColor
    ):
        self._color_rep = sRGBColor
        self._block_dict = { block: sRGBColor(*rgb, is_upscaled=True) for block, rgb in block_dict.items() }
        self._gradient_set = set(block_dict.keys())
        self._color_distance_func = ColorMetrics.rgb_distance
        self._color_distance_approx_func = None
        self._color_intermediate_func = ColorMetrics.get_intermediate_color
        
        self._block_dict = self._change_color_rep(color_rep)
        self._color_rep = color_rep
        self._update_graph()

    @classmethod   
    def from_file(cls, json_file_path):
        return cls(BlockColors.read_dict_file(json_file_path))

    @classmethod   
    def load_from_default_file(cls):
        return cls(BlockColors.read_dict_file(BlockColors._default_file))

    @staticmethod
    def read_dict_file(json_file_path):
        with open(json_file_path, 'r') as f:
            data = json.load(f)
        return {BlockType.Value(block['name']): ColorConverter.str_to_rgb(block['rgb']) for block in data}

    @staticmethod
    def generate_gradient_from_list(gradient_func, colors, steps, unique=False):
        gradient = gradient_func(colors[0],colors[1],steps,unique)
        for i in range(1,len(colors)-1):
            gradient += gradient_func(colors[i],colors[i+1],steps,unique)[1:]
        return gradient

    def _change_color_rep(self, color_rep):
        if color_rep is None:
            return { block: ColorConverter.obj_to_rgb(color) for block, color in self._block_dict.items() }
        elif self._color_rep != color_rep:
            return { block: convert_color(color, color_rep) for block, color in self._block_dict.items() }
        else:
            return self._block_dict

    def _update_graph(self):
        self._graph = nx.Graph()
        for block in self._block_dict.keys():
            self._graph.add_node(block)
        for block1, color1 in self._block_dict.items():
            for block2, color2 in self._block_dict.items():
                if block1 != block2:
                    distance = self._color_distance_func(color1, color2)
                    self._graph.add_edge(block1, block2, weight=distance)
        self._update_woking_graph()    

    def _update_woking_graph(self):
        self._working_graph = nx.subgraph_view(
            self._graph,
            filter_node= lambda node: node in self._gradient_set,
        )
        self._spanning_graph = nx.minimum_spanning_tree(self._working_graph, weight='weight')
        spanning_graph_size = sum(data['weight'] for _, _, data in self._spanning_graph.edges(data=True))
        self._mean_distance = spanning_graph_size/nx.number_of_edges(self._spanning_graph)

    def get_block_color(self, block):
        return self._block_dict.get(block)

    def get_block_distance(self, block1, block2):
        if self._graph.has_node(block1) and self._graph.has_node(block2):
            if block1 == block2:
                return 0
            return self._graph[block1][block2]['weight']
        return None

    def get_mean_distance(self):
        return self._mean_distance
    
    def get_block_dict(self, color_rep=None):
        if color_rep == self._color_rep:
            return copy.deepcopy(self._block_dict)
        else:
            return self._change_color_rep(color_rep)
    
    def get_gradient_set(self):
        return copy.deepcopy(self._gradient_set)
    
    def set_block_pallete(self, viable_blocks):
        self._gradient_set = set(self._block_dict.keys()).intersection(set(viable_blocks))
        self._update_woking_graph()

    def set_color_metrics(self, color_rep, color_distance_func: Callable, intermediate_color_func: Callable):
        update_graph = not (
            self._color_rep == color_rep and
            self._color_distance_func == color_distance_func
        )
        
        self._block_dict = self._change_color_rep(color_rep)
        self._color_rep = color_rep
        self._color_distance_func = color_distance_func
        self._color_intermediate_func = intermediate_color_func
        if update_graph:
            self._update_graph()

    def set_approximation_metric(self, color_distance_approx_func: Callable=None, approx_buffer=10):
        self._color_distance_approx_func = color_distance_approx_func
        self._approx_buffer = approx_buffer

    def get_color_distance(self, color1, color2):
        return self._color_distance_func(color1, color2)

    def closest_blocks_to_color(self, color, k=1, invalid_blocks=None):
        viable_blocks = self._gradient_set
        if invalid_blocks is not None:
            viable_blocks = viable_blocks.difference(invalid_blocks)

        if self._color_distance_approx_func is not None:
            distances = []
            for block in viable_blocks:
                block_color = self._block_dict.get(block)
                if block_color is not None:
                    distance = self._color_distance_approx_func(color, block_color)
                    distances.append(BlockColors.PathEdge(block, distance))
            sorted_distances = sorted(distances, key=lambda pe: pe.dist)
            viable_blocks = [block for block, _ in sorted_distances[:k+self._approx_buffer]]
            
        distances = []
        for block in viable_blocks:
            block_color = self._block_dict.get(block)
            if block_color is not None:
                distance = self._color_distance_func(color, block_color)
                distances.append(BlockColors.PathEdge(block, distance))

        sorted_distances = sorted(distances, key=lambda pe: pe.dist)
        closest_blocks = [block for block, _ in sorted_distances[:k]]

        return closest_blocks

    def generate_gradient(self, color1, color2, steps):
        gradient = []
        for step in range(steps):
            ratio = step / float(steps - 1)
            intermediate_color = self._color_intermediate_func(color1, color2, ratio)
            closest_block = self.closest_blocks_to_color(intermediate_color)[0]
            gradient.append(closest_block)
        return gradient

    def generate_recurcive_gradient(self, color1, color2, steps, inter_coef=0.5):
        start_block = self.closest_blocks_to_color(color1)[0]
        end_block = self.closest_blocks_to_color(color2)[0]
        
        self._block_dict.get(start_block)
        gradient = self._generate_recurcive_gradient(
            self._block_dict.get(start_block),
            self._block_dict.get(end_block),
            steps-2,
            inter_coef
        )
        gradient = [start_block]+gradient+[end_block]
        return gradient

    def _generate_recurcive_gradient(self, color1, color2, steps, inter_coef):
        if steps==0:
            return []
        
        intermediate_color = self._color_intermediate_func(color1, color2, inter_coef)
        closest_block = self.closest_blocks_to_color(intermediate_color)[0]
        intermediate_color = self._block_dict.get(closest_block)
        dist1 = self._color_distance_func(color1, intermediate_color)
        dist2 = self._color_distance_func(color2, intermediate_color)
        
        steps -= 1
        abs_dist = dist1+dist2
        ratio = dist1/abs_dist if abs_dist>0 else 0.5
        steps1 = int(steps*ratio)
        steps2 = steps-steps1
        
        gradient1 = self._generate_recurcive_gradient(color1, intermediate_color, steps1, inter_coef)
        gradient2 = self._generate_recurcive_gradient(intermediate_color, color2, steps2, inter_coef)
        
        gradient = gradient1+[closest_block]+gradient2
        return gradient

    def generate_gradient_via_graph(self, start_color, end_color, steps, dist_coef=0.5):
        
        # special case for HSV - in HSV the "shortest path" leads through middle
        # for colors that are on opposide sides of the hue spectrum. Solve through division
        # if self._color_rep == HSVColor:
        #     dist = ColorMetrics.hue_distance(start_color.hsv_h, end_color.hsv_h)
        #     if abs(dist)>(ColorMetrics._hue_max/4):
        #         num_steps +=1
        #         num_steps_l = num_steps//2
        #         num_steps_r = num_steps - num_steps_l 
        #         intermediate_color = self._color_intermediate_func(start_color, end_color, num_steps_l/num_steps)
        #         return (
        #             self.generate_gradient_via_graph( start_color, intermediate_color, num_steps_l, dist_coef) +
        #             self.generate_gradient_via_graph( intermediate_color, end_color, num_steps_r, dist_coef)[1:]
        #         )

        start_block = self.closest_blocks_to_color(start_color)[0]
        end_block = self.closest_blocks_to_color(end_color, invalid_blocks=[start_block])[0] # better invalid block?
        
        distance = self._graph[start_block][end_block]["weight"]/steps
        distance += distance*dist_coef
        # In CIELAB (delta) distances to White are huge and they screw up distance
        
        graph = nx.subgraph_view(
            self._working_graph,
            filter_edge=lambda node1, node2: 
                self._graph[node1][node2]['weight'] < distance or
                self._spanning_graph.has_edge(node1, node2)
        )
        
        path = nx.shortest_path(
            graph, source=start_block, target=end_block, weight='weight'
        )

        distances = []
        path_dlist = dllist()

        for i in range(len(path)-1):
            block_from, block_to = path[i], path[i+1]
            distance = self._graph[block_from][block_to]["weight"]
            path_node = path_dlist.append(BlockColors.PathEdge(path[i], distance))
            distances.append(BlockColors.PathEdge(path_node, distance))    
        path_dlist.append(BlockColors.PathEdge(path[-1],0))
        
        sorted_distances = sorted(distances, key=lambda pe: pe.dist)
        
        while path_dlist.size > max(steps,2):
            path_node, _ = sorted_distances.pop(0)
            if path_node.owner is not None and path_node.value.dist <= sorted_distances[0].dist:
                left, right = path_node, path_node.next
                if left.prev is None:
                    node_to_del = right
                elif right.next is None:
                    node_to_del = left
                else:
                    node_to_del = left if left.value.dist < right.value.dist else right

                node = node_to_del.prev
                path_dlist.remove(node_to_del)
                
                new_dist = self._graph[node.value.node][node.next.value.node]["weight"]
                node.value = BlockColors.PathEdge(node.value.node, new_dist)
                index = bisect_left(sorted_distances, new_dist, key=lambda pe: pe.dist)
                sorted_distances.insert(index, BlockColors.PathEdge(node, new_dist))

        # # Enlarge to steps size
        # path = []
        # if path_dlist.size <steps:
        #     full_dist = sum([dist for _,dist in path_dlist.itervalues()][:-1])
        #     step_dist = full_dist/steps
        #     current_dist = 0
        #     for block, dist in path_dlist.itervalues():
        #         num = (current_dist + dist/2)/step_dist
        #         num = math.floor(num)
        #         for _ in range(len(path),num):
        #             path.append(block)
        #         current_dist +=dist
        # else:
        #     path = [block  for block,_ in path_dlist.itervalues()]
        
        # return path
        
        return [block  for block,_ in path_dlist.itervalues()]

    # region depricated
    # def _update_neighbors(self):
    #     negihbours dist sorted
    #     self._closest_neighbors = {}
    #     for block1, color1 in self._block_dict.items():
    #         distances = []
    #         for block2, color2 in self._block_dict.items():
    #             if block1 != block2:
    #                 distances.append(self._graph[block1][block2]['weight'])
    #         sorted_distances = sorted(distances, key=lambda x: x[1])    
    #         self._closest_neighbours[block1] = [block for block,_ in sorted_distances]
    
    # def generate_gradient_via_graph_(self, start_color, end_color, num_steps):
    #     start_block = self.closest_blocks_to_color(start_color)[0]
    #     end_block = self.closest_blocks_to_color(end_color, invalid_blocks=[start_block])[0] # better invalid block?
        
    #     path = [start_block, end_block]
    #     while len(path)<num_steps:
    #         distances = []
    #         for i in range(len(path) - 1):
    #             distance = self._graph[path[i]][path[i + 1]]["weight"] 
    #             distances.append(BlockColors.PathEdge(i, distance))

    #         # Sort in descending order based on distance
    #         sorted_distances = sorted(distances, key=lambda pe: pe.dist, reverse=True)
    #         for i, distance in sorted_distances:
    #             graph = nx.subgraph_view(
    #                 self._working_graph,
    #                 filter_edge=lambda node1, node2: self._graph[node1][node2]['weight'] < distance
    #             )
    #             if nx.has_path(graph, source=path[i], target=path[i + 1]):
    #                 sub_path = nx.shortest_path(
    #                     graph, source=path[i], target=path[i + 1], weight='weight'
    #                 )
    #                 path = path[:i+1] + sub_path[1:-1] + path[i+1:]
    #                 break
        
    #     return path
    
    # endregion