import os
import torch
import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from .architectures import AuxModel, BiomassModel
from ..core.config import settings

class ModelManager:
    """Singleton pattern or persistent instance to manage models and memory out of route logic"""
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.aux_models = []
        self.main_models = []
        
        self.transform = A.Compose([
            A.Resize(settings.img_size, settings.img_size),
            A.Normalize(),
            ToTensorV2()
        ])
        self.is_loaded = False

    def load_models(self):
        if self.is_loaded:
            return
            
        print(f"Loading models to {self.device}...")
        
        # Load Aux Models
        for seed in settings.aux_seeds:
            for fold in range(5):
                path = os.path.join(settings.aux_model_dir, f'best_aux_only_seed{seed}_fold{fold}.pth')
                if not os.path.exists(path):
                    path_alt = os.path.join(settings.aux_model_dir, f'best_aux_seed{seed}_fold{fold}.pth')
                    if os.path.exists(path_alt):
                        path = path_alt
                    else: continue
                
                ckpt = torch.load(path, map_location=self.device, weights_only=False)
                model = AuxModel(settings.model_name).to(self.device).eval()
                model.load_state_dict(ckpt['model_state_dict'])
                self.aux_models.append((model, ckpt))

        # Load Main Models
        for seed in settings.seeds:
            for fold in settings.fold_weights.keys():
                path = os.path.join(settings.main_model_dir, f'best_model_seed{seed}_fold{fold}.pth')
                if not os.path.exists(path): continue
                
                ckpt = torch.load(path, map_location=self.device, weights_only=False)
                model = BiomassModel(settings.model_name, img_size=settings.img_size).to(self.device).eval()
                model.load_state_dict(ckpt['model_state_dict'] if 'model_state_dict' in ckpt else ckpt)
                self.main_models.append((model, ckpt, settings.fold_weights[fold]))
                
        self.is_loaded = True
        print(f"Loaded {len(self.aux_models)} Aux Models and {len(self.main_models)} Main Models successfully.")

    def predict(self, image_bytes: bytes) -> dict:
        if not self.is_loaded:
            raise RuntimeError("Models are not loaded.")

        # Decode Image
        nparr = np.frombuffer(image_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Could not decode image.")
            
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # Transform
        tensor_img = self.transform(image=img)['image'].unsqueeze(0).to(self.device)

        with torch.no_grad():
            # ---- Stage 1: Aux Prediction ----
            aux_preds_accum = 0
            for model, ckpt in self.aux_models:
                res = model(tensor_img).cpu().numpy()
                if 'tab_scaler' in ckpt: 
                    res = ckpt['tab_scaler'].inverse_transform(res)
                aux_preds_accum += res

            # Average aux predictions
            predicted_tabular = aux_preds_accum / max(1, len(self.aux_models))

            # ---- Stage 2: Main Prediction ----
            final_biomass_accum = 0
            total_w = 0
            for model, ckpt, weight in self.main_models:
                tab_input = predicted_tabular.copy()
                if ckpt.get('tabular_scaler'): 
                    tab_input = ckpt['tabular_scaler'].transform(tab_input)
                
                tab_tensor = torch.tensor(tab_input, dtype=torch.float32).to(self.device)
                
                res = model(tensor_img, tab_tensor).cpu().numpy()
                if ckpt.get('target_scaler'): 
                    res = ckpt['target_scaler'].inverse_transform(res)
                
                final_biomass_accum += (res * weight)
                total_w += weight

        if total_w == 0:
            raise RuntimeError("No Main Models available for prediction.")

        final_preds = final_biomass_accum / total_w
        
        # Format results (equivalent to np.maximum(0, val))
        results = {}
        for j, target in enumerate(settings.targets):
            results[target] = max(0.0, float(final_preds[0, j]))
            
        return results

model_manager = ModelManager()