import os
import glob
import json
import cv2
import numpy as np
from fis_common_function import FIS_common_acquisition as FIS


class CreationPatterns:
    """
    Universal ONE-PIX pattern loader.
    Loads patterns from a folder specified in ima_custom_config.json.
    """

    def __init__(self, height=0, width=0):
        # --- Structure standard ONE-PIX ---
        self.height = height
        self.width = width
        self.acquisition_results = {}
        self.patterns = []
        self.patterns_order = []
        self.nb_patterns = 0
        self.interp_method = cv2.INTER_AREA

        # --- Lecture du JSON de config ---
        base_path = os.path.dirname(os.path.abspath(__file__))
        self.config_path = os.path.join(base_path, "custom_config.json")

        if not os.path.exists(self.config_path):
            raise FileNotFoundError(
                f"❌ Configuration file not found: {self.config_path}\n"
                f"Please create it with a 'folder_path' key."
            )

        with open(self.config_path, "r") as f:
            config = json.load(f)

        self.folder_path = config.get("folder_path", "").strip()
        if not self.folder_path or not os.path.isdir(self.folder_path):
            raise ValueError(f"❌ Invalid or missing folder_path in {self.config_path}: '{self.folder_path}'")

        print(f"📁 Using patterns from: {self.folder_path}")

    # ------------------------------------------------------------------
    def _normalize_pattern(self, img):
        """Ensure all patterns are float32 in [0,1]."""
        if img is None:
            return None
        img = img.astype(np.float32)
        if img.max() > 1.5:
            img = img / 255.0
        return np.clip(img, 0.0, 1.0)

    # ------------------------------------------------------------------
    def _load_all_patterns(self, folder):
        """Loads every supported file in the folder (jpg, png, npy, npz)."""
        exts = ["*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tif", "*.tiff", "*.npy", "*.npz"]
        all_files = []
        for ext in exts:
            all_files.extend(glob.glob(os.path.join(folder, ext)))

        if not all_files:
            raise FileNotFoundError(f"No valid pattern files found in '{folder}'")

        all_files = sorted(all_files, key=os.path.getmtime)

        for filepath in all_files:
            name = os.path.splitext(os.path.basename(filepath))[0]
            ext = os.path.splitext(filepath)[1].lower()

            try:
                if ext in [".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"]:
                    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
                    img = self._normalize_pattern(img)
                    if img is not None:
                        self.patterns.append(img)
                        self.patterns_order.append(name)

                elif ext == ".npy":
                    arr = np.load(filepath)
                    if arr.ndim == 3:
                        for i in range(arr.shape[0]):
                            self.patterns.append(self._normalize_pattern(arr[i]))
                            self.patterns_order.append(f"{name}_{i}")
                    else:
                        self.patterns.append(self._normalize_pattern(arr))
                        self.patterns_order.append(name)

                elif ext == ".npz":
                    data = np.load(filepath)
                    for key in data:
                        arr = np.float32(data[key])
                        self.patterns.append(self._normalize_pattern(arr))
                        self.patterns_order.append(f"{name}_{key}")

            except Exception as e:
                print(f"⚠️ Failed to load {filepath}: {e}")

        if not self.patterns:
            raise RuntimeError(f"No valid patterns could be loaded from {folder}")

    # ------------------------------------------------------------------
    def creation_patterns(self):
        """Main method for ImagingMethodBridge."""
        self._load_all_patterns(self.folder_path)

        ref_shape = self.patterns[0].shape
        cleaned_patterns = []
        for p in self.patterns:
            if p.shape != ref_shape:
                p = cv2.resize(p, (ref_shape[1], ref_shape[0]), interpolation=cv2.INTER_AREA)
            cleaned_patterns.append(p)

        self.patterns = [np.float32(p) for p in cleaned_patterns]
        self.nb_patterns = len(self.patterns)
        self.acquisition_results["patterns"] = self.patterns
        self.acquisition_results["patterns_order"] = self.patterns_order

        print(f"✅ Loaded {self.nb_patterns} patterns from '{self.folder_path}'")
        print(f"→ Shape: {self.patterns[0].shape}, dtype: {self.patterns[0].dtype}, range: [{self.patterns[0].min()}, {self.patterns[0].max()}]")
        return self.patterns

    # ------------------------------------------------------------------
    def sequence_order(self):
        """Return pattern order (standard ONE-PIX API)."""
        if not self.patterns_order:
            self.patterns_order = [f"pattern_{i}" for i in range(self.nb_patterns)]
        self.acquisition_results["patterns_order"] = self.patterns_order
        return self.patterns_order

    # ------------------------------------------------------------------
    def save_raw_data(self, acquisition_class, path=None):
        saver = FIS.FisCommonAcquisition(acquisition_class)
        saver.save_raw_data(path)
