import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
from onepix import utils as utils
from pathlib import Path
from datetime import date
import time
import orjson

class Reconstruction:

    def __init__(self,acquisition_dict):
        self.reconstruction_results={}

        self.reconstruction_results["spectra"] = np.asarray(acquisition_dict["spectra"])
        self.reconstruction_results["wavelengths"]=np.asarray(acquisition_dict["wavelengths"])
        self.reconstruction_results["cluster_name"] = acquisition_dict["patterns_order"]
        
        patterns_list = acquisition_dict["patterns"]
        masks = [np.array(mask) for mask in patterns_list]
        self.reconstruction_results["masks"] = masks

        self.reconstruction_results["rgb"]=np.asarray(acquisition_dict["rgb"])


    def image_reconstruction(self):
        pass

    def save_reconstructed_image(self, header, filename, save_path=None):
        save_path = Path(__file__).resolve().parent.parent.parent.parent / "measure"
        if not os.path.isdir(save_path):
            os.mkdir(save_path)

        fdate = date.today().strftime("%d_%m_%Y")
        actual_time = time.strftime("%H-%M-%S")
        folder_name = f"ONE-PIX_HAS_reconstruction_{fdate}_{actual_time}"
        os.mkdir(os.path.join(save_path,folder_name))

        results_filename = f"reconstruction_results_{fdate}_{actual_time}.json"
        with open(os.path.join(save_path, folder_name, results_filename), "wb") as f:
            f.write(
                orjson.dumps(
                    self.reconstruction_results,
                    option=orjson.OPT_SERIALIZE_NUMPY, 
                    default=utils.default
                )
            )


    def get_result_to_plot(self):
        self.reconstruction_results["result2plot"]=self.build_overlay_with_spectra()
        return self.reconstruction_results["result2plot"]
    

    def build_overlay_with_spectra(self, alpha=0.5, colormap_fn=None, figsize=(6, 4)):
        """
        Construit une image unique combinant :
          - l'image RGB + masques colorés
          - un graphe matplotlib des spectres
        """

        # ------------------------------
        # Chargement + Debug RGB
        # ------------------------------
        rgb = self.reconstruction_results["rgb"]
        print("\n=== DEBUG RGB =========================")
        print("Raw RGB SHAPE:", rgb.shape, "DTYPE:", rgb.dtype)

        rgb = np.asarray(rgb)

        # Grayscale → RGB
        if len(rgb.shape) == 2:
            print("RGB is grayscale → converting to BGR")
            rgb = cv2.cvtColor(rgb, cv2.COLOR_GRAY2BGR)

        # RGBA → RGB
        elif rgb.shape[2] == 4:
            print("RGB has 4 channels → converting BGRA→BGR")
            rgb = cv2.cvtColor(rgb, cv2.COLOR_BGRA2BGR)

        rgb = rgb.astype(np.uint8)

        h, w = rgb.shape[:2]
        print("Final RGB SHAPE:", rgb.shape)

        # ------------------------------
        # Debug des masques
        # ------------------------------
        masks = self.reconstruction_results["masks"]
        print("\n=== DEBUG MASKS =======================")
        fixed_masks = []
        for idx, mask in enumerate(masks):
            mask = np.array(mask)
            print(f"Mask {idx} original shape:", mask.shape, "dtype:", mask.dtype)

            # Convert to 255 mask
            if mask.dtype != np.uint8:
                mask = (mask > 0).astype(np.uint8) * 255

            # Resize if needed
            if mask.shape != (h, w):
                print(f" → resizing mask {idx} from {mask.shape} to {(h, w)}")
                mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)

            print(f"Mask {idx} final shape:", mask.shape)
            fixed_masks.append(mask)

        masks = fixed_masks

        # ------------------------------
        # Création overlay + debug
        # ------------------------------
        print("\n=== DEBUG OVERLAY =====================")
        overlay = np.zeros((h, w, 3), dtype=np.uint8)

        # Colors
        n = len(masks)
        if colormap_fn:
            colors = colormap_fn(n+1)[1:]
            colors = [tuple(np.array(c) / 255) for c in colors]
        else:
            cmap = plt.get_cmap("tab10")
            colors = [cmap(i % 10)[:3] for i in range(n)]

        for i, mask in enumerate(masks):
            color = (np.array(colors[i]) * 255).astype(np.uint8)
            overlay[mask > 0] = color

        print("Overlay SHAPE:", overlay.shape, "DTYPE:", overlay.dtype)

        # ------------------------------
        # Vérification finale avant addWeighted
        # ------------------------------
        print("\n=== DEBUG BEFORE ADDWEIGHTED ==========")
        print("RGB shape :", rgb.shape)
        print("OVL shape :", overlay.shape)

        if rgb.shape != overlay.shape:
            print("!! SHAPES MISMATCH → resizing overlay")
            overlay = cv2.resize(overlay, (rgb.shape[1], rgb.shape[0]))

        print("FINAL OVL shape:", overlay.shape)

        blended = cv2.addWeighted(rgb, 1 - alpha, overlay, alpha, 0)
        print("Blended SHAPE:", blended.shape)

        # ------------------------------
        # Graphe matplotlib
        # ------------------------------
        spectra = self.reconstruction_results["spectra"]
        wavelengths = self.reconstruction_results["wavelengths"]

        fig, ax = plt.subplots(figsize=figsize, dpi=100)
        for i in range(n):
            ax.plot(wavelengths, spectra[i], color=colors[i], alpha=0.9, label=f"Cluster {i+1}")
        ax.set_xlabel("Wavelength (nm)")
        ax.set_ylabel("Intensity")
        ax.legend()
        ax.grid(True, linestyle="--", alpha=0.3)
        fig.tight_layout()

        fig.canvas.draw()
        graph_img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        graph_img = graph_img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        plt.close(fig)

        print("\n=== DEBUG GRAPH ======================")
        print("Graph image SHAPE:", graph_img.shape)

        graph_img = cv2.resize(graph_img, (w, h))
        print("Graph resized SHAPE:", graph_img.shape)

        # ------------------------------
        # Concat final
        # ------------------------------
        combined = np.concatenate([blended, graph_img], axis=1)
        print("\n=== FINAL IMAGE =======================")
        print("Combined SHAPE:", combined.shape)
        print("=======================================\n")

        return combined
