# models/regression.py

import pandas as pd
import numpy as np
from typing import Literal, Dict, Any
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import MinMaxScaler

RegressionModelType = Literal['linear', 'xgb', 'mlp']

class LoadLimitRegressor:
    """
    Regressor สำหรับหา load limit จาก temperature
    พยากรณ์โหลดสูงสุดที่เป็นไปได้จากค่าความร้อน
    """
    def __init__(self, model_type: RegressionModelType = 'xgb'):
        self.model_type = model_type
        self.scaler_X = MinMaxScaler()
        self.scaler_y = MinMaxScaler()
        self.model = None
        self.feature_names = None

    def fit(self, df: pd.DataFrame, temp_col: str = 'MaxTemp', load_col: str = 'mw_max'):
        """
        เทรน regression model เพื่อหา relationship ระหว่างอุณหภูมิ ➜ โหลดสูงสุด
        """
        X = df[[temp_col]].values
        y = df[[load_col]].values

        X_scaled = self.scaler_X.fit_transform(X)
        y_scaled = self.scaler_y.fit_transform(y)

        if self.model_type == 'linear':
            self.model = LinearRegression()
        elif self.model_type == 'xgb':
            self.model = GradientBoostingRegressor(n_estimators=100, learning_rate=0.1, max_depth=3)
        elif self.model_type == 'mlp':
            self.model = MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=1000)
        else:
            raise ValueError(f"Unsupported model type: {self.model_type}")

        self.model.fit(X_scaled, y_scaled.ravel())
        self.feature_names = [temp_col]

    def predict(self, max_temp_values: np.ndarray) -> np.ndarray:
        """
        พยากรณ์โหลดสูงสุดที่เป็นไปได้ จากค่าความร้อน
        """
        X_scaled = self.scaler_X.transform(max_temp_values.reshape(-1, 1))
        y_scaled = self.model.predict(X_scaled).reshape(-1, 1)
        y = self.scaler_y.inverse_transform(y_scaled)
        return y.flatten()

    def evaluate(self, df: pd.DataFrame, temp_col: str = 'MaxTemp', load_col: str = 'mw_max') -> Dict[str, float]:
        """
        ประเมินผลโมเดล regression
        """
        y_true = df[load_col].values
        y_pred = self.predict(df[temp_col].values)

        return {
            'rmse': np.sqrt(mean_squared_error(y_true, y_pred)),
            'mae': mean_absolute_error(y_true, y_pred),
            'r2': r2_score(y_true, y_pred)
        }


class CapacityAnalyzer:
    """
    วิเคราะห์ capacity จริง (Real Capacity) เทียบกับทฤษฏี 80%
    
    ใช้ features เพิ่มเติม (TEPR, STRR, temp margin) เพื่อประมาณ capacity จริง
    ที่อาจสูงกว่า 80% theoretical
    """
    def __init__(self, model_type: RegressionModelType = 'xgb'):
        self.model_type = model_type
        self.scalers = {}
        self.model = None
        self.feature_names = None

    def fit(self, df: pd.DataFrame, 
            target_col: str = 'mw_max',
            feature_cols: list = None,
            calibration_factor: float = 1.0) -> None:
        """
        เทรน model สำหรับหา real capacity
        
        Args:
            df: DataFrame ที่มีข้อมูล load, temp, measurement
            target_col: ชื่อ column เป้าหมาย (load actual)
            feature_cols: ชื่อ columns ที่ใช้เป็น features
                         ถ้า None ใช้ default: ['MaxTemp', 'temp_margin_available', 'tepr_mean', 'strr_mean']
            calibration_factor: factor สำหรับ calibrate ผล (default 1.0)
        """
        if feature_cols is None:
            # Default features
            available_cols = df.columns.tolist()
            feature_cols = []
            if 'MaxTemp' in available_cols:
                feature_cols.append('MaxTemp')
            if 'temp_margin_available' in available_cols:
                feature_cols.append('temp_margin_available')
            if 'tepr_mean' in available_cols:
                feature_cols.append('tepr_mean')
            if 'strr_mean' in available_cols:
                feature_cols.append('strr_mean')
            
            # ถ้าไม่มี feature พอ ใช้เฉพาะ MaxTemp
            if not feature_cols:
                feature_cols = ['MaxTemp']
        
        # ตรวจสอบว่ามี feature ทั้งหมดหรือไม่
        available_features = [col for col in feature_cols if col in df.columns]
        if not available_features:
            raise ValueError(f"❌ ไม่มี features {feature_cols} ใน data!")
        
        X = df[available_features].values
        y = df[[target_col]].values
        
        # Normalize each feature
        self.scalers = {}
        X_scaled = np.zeros_like(X)
        for i, col in enumerate(available_features):
            scaler = MinMaxScaler()
            X_scaled[:, i] = scaler.fit_transform(X[:, i:i+1]).flatten()
            self.scalers[col] = scaler
        
        # Train model
        if self.model_type == 'linear':
            self.model = LinearRegression()
        elif self.model_type == 'xgb':
            self.model = GradientBoostingRegressor(n_estimators=100, learning_rate=0.1, max_depth=3)
        elif self.model_type == 'mlp':
            self.model = MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=1000)
        else:
            raise ValueError(f"Unsupported model type: {self.model_type}")
        
        # Normalize target
        self.scaler_y = MinMaxScaler()
        y_scaled = self.scaler_y.fit_transform(y)
        
        self.model.fit(X_scaled, y_scaled.ravel())
        self.feature_names = available_features
        self.calibration_factor = calibration_factor

    def predict_capacity(self, df: pd.DataFrame) -> np.ndarray:
        """
        พยากรณ์ real capacity จริง
        """
        if self.model is None:
            raise ValueError("❌ Model ยังไม่ได้ train! ให้เรียก fit() ก่อน")
        
        # ดึงแต่ features ที่มี
        X = df[self.feature_names].values
        
        # Scale features
        X_scaled = np.zeros_like(X)
        for i, col in enumerate(self.feature_names):
            if col in self.scalers:
                X_scaled[:, i] = self.scalers[col].transform(X[:, i:i+1]).flatten()
        
        # Predict
        y_scaled = self.model.predict(X_scaled).reshape(-1, 1)
        y = self.scaler_y.inverse_transform(y_scaled)
        
        # Apply calibration factor
        return (y.flatten() * self.calibration_factor)

    def evaluate(self, df: pd.DataFrame, 
                actual_col: str = 'mw_max') -> Dict[str, Any]:
        """
        ประเมินผลโมเดล capacity analysis
        """
        y_true = df[actual_col].values
        y_pred = self.predict_capacity(df)
        
        metrics = {
            'rmse': np.sqrt(mean_squared_error(y_true, y_pred)),
            'mae': mean_absolute_error(y_true, y_pred),
            'r2': r2_score(y_true, y_pred),
            'features_used': self.feature_names
        }
        
        return metrics