import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import json
import grpc
import minecraft_pb2_grpc
import numpy as np
import test_blocks_showcase as bs
import matplotlib.pyplot as plt
from minecraft_pb2 import *
from utils_coordinates import *
from utils_blocks_categories import *
from PIL import ImageColor
from collections import defaultdict

WORKING_DIR = os.path.join("py_scripts","block_color_extraction")
BLOCK_DATA_FILE = 'blocks_1.20.2.json'
OUTPUT_FILE = 'block_to_color.json'
OUTPUT_NAME_FILE = 'block_to_name.json'

def load_block_data(json_file):
    with open(os.path.join(WORKING_DIR,json_file), 'r') as f:
        data = json.load(f)
    block_dict = {block['name'].upper(): ImageColor.getrgb(block['rgb']) for block in data}
    return block_dict

class NameConverter:

    _prefixes = tuple(sorted(
        ("WAXED_","INFESTED_","CHIPPED_","DAMAGED_"), #"SUSPICIOUS_"
        key=len, reverse=True
    ))
    _suffixes = tuple(sorted(
        ("_BLOCK","_SLAB","_STAIRS","_FENCE","_FENCE_GATE","_WALL","_PANE","_BUTTON","_SIGN","_BANNER",
        "_WALL_BANNER","_PRESSURE_PLATE","_HANGING_SIGN","_WALL_HANGING_SIGN"),
        key=len, reverse=True
    ))
    
    def __init__(self, block_dict:dict) -> None:
        self.block_dict = block_dict
    
    @staticmethod
    def _average_color(rgb_colors):
        if not rgb_colors:
            return None

        # Initialize sums
        avg_color = np.mean(rgb_colors, axis=0)
        tuple(avg_color.astype(int))

        # Return the average color
        return tuple(avg_color.astype(int))

    def _try_numbers(self, name):
        result_sides = []
        for i in range(8):
            result = self.block_dict.get(name+str(i))
            if result is not None:
                result_sides.append(result)
        return NameConverter._average_color(result_sides)

    def _get_color(self, name:str, d=0):

        result = self.block_dict.get(name,
            self.block_dict.get(name+"S",
                self.block_dict.get(name+"_SIDE",
                    self.block_dict.get(name+"_SIDES")
                )
            )
        )
        if result is None:
            if name.endswith(self._suffixes):
                for suffix in self._suffixes:
                    if name.endswith(suffix):
                        new_name = name.removesuffix(suffix)
                        result = self._get_color(new_name,d)
                        break
            elif name.startswith(self._prefixes):
                for prefix in self._prefixes:
                    if name.startswith(prefix):
                        new_name = name.removeprefix(prefix)
                        result = self._get_color(new_name,d)
                        break
            elif name in ["DISPENSER","DROPPER"]:
                result = self._get_color("FURNACE")
            elif name.startswith("SMOOTH_"):
                new_name = name[name.find('_')+1:]
                result = self.block_dict.get(new_name+"_TOP",
                    self.block_dict.get(new_name+"_BLOCK_TOP")
                )
            elif name.endswith("_DOOR"):
                result = self.block_dict.get(name+"_BOTTOM")
            elif name.endswith("_WOOD"):
                new_name = name.replace("_WOOD","_LOG")
                result = self.block_dict.get(new_name)
            elif name.endswith("_CARPET"):
                new_name = name.replace("_CARPET","_WOOL")
                result = self.block_dict.get(new_name)
            elif name.endswith("_HYPHAE"):
                new_name = name.replace("_HYPHAE","_STEM")
                # new_name = new_name.removeprefix("STRIPPED_")
                result = self.block_dict.get(new_name)
            elif name.startswith("POTTED_"):
                pot = "FLOWER_POT"
                new_name = name.removeprefix("POTTED_")
                colors = [
                    self.block_dict.get(pot),
                    self._get_color(new_name)
                ]
                if colors[1] is not None:
                    result = NameConverter._average_color(colors)
            elif name.find("_WALL_")>=0:
                new_name = name.replace("_WALL_","_")
                result = self.block_dict.get(new_name)
            elif d==0:
                for suffix in ["_BLOCK","_BRICK","_PLANK","_CONCRETE"]:
                    result = self._get_color(name+suffix,1)
                    if result != None:
                        break
            else:
                for suffix in ["_","_SIDE","_STAGE",]:
                    result = self._try_numbers(name+suffix)
                    if result != None:
                        break
        return result
    
    def get_color(self, name:str):
        return self._get_color(name)

class TrackingDict(dict):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.accessed_keys = set()

    def get(self, key, default=None):
        # Store the accessed key in the instance's set
        self.accessed_keys.add(key)
        # Call the original dict's get method
        return super().get(key, default)

    def get_accessed_keys(self):
        # Return the list of accessed keys
        return list(self.accessed_keys)
    
    def get_non_accessed_keys(self):
        # Calculate the set of non-accessed keys
        non_accessed_keys = set(self.keys()) - self.accessed_keys
        return list(non_accessed_keys)
          
def test_conversion(blocks, block_dict, folder = "block_pairs"):
    paired_blocks = []
    unpaired_blocks = []
    
    tracking_block_dict = TrackingDict(block_dict)
    nc = NameConverter(tracking_block_dict)
    
    for block_enum in blocks:
        block_name = BlockType.Name(block_enum)
        if nc.get_color(block_name) is not None:
            paired_blocks.append(block_enum)
        else:
            unpaired_blocks.append(block_enum)
        
    os.makedirs(os.path.join(WORKING_DIR, folder), exist_ok=True)
    # Write results to separate files
    with open(os.path.join(WORKING_DIR, folder,'paired_blocks.txt'), 'w') as paired_file:
        paired_file.write("Paired blocks:\n")
        for paired in paired_blocks:
            block_name = BlockType.Name(paired)
            paired_file.write(f"{block_name} : {nc.get_color(block_name)}" + ",\n")

    with open(os.path.join(WORKING_DIR, folder,'unpaired_blocks.txt'), 'w') as unpaired_file:
        unpaired_file.write("Unpaired blocks (in BlockType):\n")
        for unpaired in unpaired_blocks:
            block_name = BlockType.Name(unpaired)
            unpaired_file.write(block_name + ",\n")

    with open(os.path.join(WORKING_DIR, folder,'dict_unpaired_blocks.txt'), 'w') as dict_unpaired_file:
        dict_unpaired_file.write("Unpaired blocks (in block_dict):\n")
        for dict_unpaired in sorted(tracking_block_dict.get_non_accessed_keys()):
            dict_unpaired_file.write(dict_unpaired + ",\n")
            
    with open(os.path.join(WORKING_DIR, folder,'dict_all_blocks.txt'), 'w') as dict_all_file:
        dict_all_file.write("Blocks (in block_dict):\n")
        for dict_block in reversed(block_dict):
            dict_all_file.write(dict_block + ",\n")

    print("Results written to paired_blocks.txt, unpaired_blocks.txt, and dict_unpaired_blocks.txt")
    
    try:
        client = minecraft_pb2_grpc.MinecraftServiceStub(grpc.insecure_channel('localhost:5001'))
        bshow = bs.BlockShowcase(client)
        bshow.place_blocks_in_all_orientations(unpaired_blocks, Coordinates(6969,0,0), description="unpaired")
    except:
        print("Server is not running. Skiping placing unpaired blocks")

def plot_color_overlap(blocks, block_dict):
    nc = NameConverter(block_dict)
    colors = []
    for block_enum in blocks:
            block_name = BlockType.Name(block_enum)
            block_color = nc.get_color(block_name)
            if block_color is not None:
                colors.append((block_name,block_color))
        
    color_strings = defaultdict(list)

    # Populate the defaultdict
    for string, color in colors:
        hex_color = '#{0:02X}{1:02X}{2:02X}'.format(*color)
        color_strings[hex_color].append(string)

    # Print strings with the same color
    for color, strings in color_strings.items():
        if len(strings) > 1:  # Only print if there are multiple strings with the same color
            print(f"Strings with color {color}: {', '.join(strings)}")
            
    color_counts = {color: len(strings) for color, strings in color_strings.items()}

    # Sort colors by their counts
    sorted_colors = sorted(color_counts.items(), key=lambda x: x[1], reverse=True)

    # Extract colors and counts for plotting
    colors = [color for color, count in sorted_colors]
    counts = [count for color, count in sorted_colors]

    # Plotting
    plt.figure(figsize=(10, 6))
    plt.bar(np.arange(len(colors)), counts, color=colors)
    plt.xticks(np.arange(len(colors)), ["" for _ in range(len(colors))])  # Removing x-axis ticks
    plt.ylabel('Number of Strings')
    plt.title('Number of Strings by Color')
    plt.show()

def enum_name_convert(block_dict):
    nc = NameConverter(block_dict)
    reversed_dict = {value: key.lower() for key, value in block_dict.items()}
    
    block_data = []
    for block_enum in BlockType.values():
            block_name = BlockType.Name(block_enum)
            block_color = nc.get_color(block_name)
            if reversed_dict.get(block_color) is not None:
                block_data.append({
                    'enum': block_enum,
                    'file_name': reversed_dict[block_color] +'.png'
                })
    with open(os.path.join(WORKING_DIR, OUTPUT_NAME_FILE), 'w') as f:
        json.dump(block_data, f, indent=4)
        
    print(f"Results written to {OUTPUT_NAME_FILE}")

def enum_convert(block_dict):
    nc = NameConverter(block_dict)
    
    block_data = []
    for block_enum in BlockType.values():
            block_name = BlockType.Name(block_enum)
            block_color = nc.get_color(block_name)
            if block_color is not None:
                block_data.append({
                    'name': block_name,
                    'rgb': '#{0:02X}{1:02X}{2:02X}'.format(*block_color)
                })
    with open(os.path.join(WORKING_DIR, OUTPUT_FILE), 'w') as f:
        json.dump(block_data, f, indent=4)
        
    print(f"Results written to {OUTPUT_FILE}")

block_dict = load_block_data(BLOCK_DATA_FILE)
# test_conversion(BlockType.values(), block_dict)
# plot_color_overlap(BlockType.values(), block_dict)
enum_name_convert(block_dict)
# enum_convert(block_dict)
