app/services/ml_logic.py

from .ml_models import ml_models
from .scoring import calculate_vague_score, calculate_concrete_score, analyze_sentiment
import re
import joblib
import os
import pandas as pd
import numpy as np

# Path configurations
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODEL_DIR = os.path.join(BASE_DIR, "ml_models")

ENSEMBLE_PATH = os.path.join(MODEL_DIR, 'ensemble_model.pkl')
FEATURE_COLS_PATH = os.path.join(MODEL_DIR, 'all_feature_columns.pkl')
CAT_MAPPING_PATH = os.path.join(MODEL_DIR, 'category_to_greenwashing_mapping.pkl')
BINARY_MAPPING_PATH = os.path.join(MODEL_DIR, 'binary_to_report_name_mapping.pkl')

_ensemble_model = None
_feature_cols = None
_binary_mapping = None

def load_artifacts():
    global _ensemble_model, _feature_cols, _binary_mapping
    
    if _ensemble_model and _feature_cols:
        return _ensemble_model, _feature_cols, _binary_mapping

    try:
        if os.path.exists(ENSEMBLE_PATH):
            print(f"[ML] Loading Ensemble Model from {ENSEMBLE_PATH}...")
            _ensemble_model = joblib.load(ENSEMBLE_PATH)
            _feature_cols = joblib.load(FEATURE_COLS_PATH)
            
            if os.path.exists(BINARY_MAPPING_PATH):
                _binary_mapping = joblib.load(BINARY_MAPPING_PATH)
            else:
                # Fallback mapping if file missing
                _binary_mapping = {0: 'Not Greenwashing (Low)', 1: 'Greenwashing (High/Medium)'}
                
            print(f"[ML] Ensemble Model Loaded. Features: {_feature_cols}")
            return _ensemble_model, _feature_cols, _binary_mapping
    except Exception as e:
        print(f"[ML] Failed to load artifacts: {e}")
    
    return None, None, None

def train_model(data: list[dict]):
    """
    Legacy training function kept for compatibility but effectively disabled
    as we are now using the pre-trained Ensemble Model.
    """
    print("[ML] Train requested, but system is now using pre-trained Ensemble Model.")
    return 0.0

def predict_greenwashing_risk(text, company_name="Unknown", features_dict=None):
    """
    Predict greenwashing risk using Ensemble Model if features are provided.
    Fallback to heuristic if only text is available.
    """
    model, features, binary_map = load_artifacts()
    
    # 1. Prediction using Ensemble Model (Feature-based)
    if model and features and features_dict:
        try:
            # Prepare input dataframe with correct column order
            input_data = {}
            for col in features:
                # Handle typo in specific user column "frequecy"
                val = features_dict.get(col)
                if val is None:
                    # Fallback for known variations
                    if col == 'Green Keyword frequecy': 
                         val = features_dict.get('Green Keyword Frequency', 0)
                    elif col == 'Emission Sentiment ': # Note space
                         val = features_dict.get('Emission Sentiment', 0)
                    else:
                         val = 0
                input_data[col] = [float(val)]
            
            df = pd.DataFrame(input_data)
            
            # Predict
            pred_binary = model.predict(df)[0]
            pred_proba = model.predict_proba(df)[0] # [prob_0, prob_1]
            prob_gw = pred_proba[1]
            
            # granular mapping based on probability
            if prob_gw >= 0.75:
                risk_label = "High"
                label_text = "High Risk"
            elif prob_gw >= 0.35:
                risk_label = "Medium"
                label_text = "Medium Risk"
            else:
                risk_label = "Low"
                label_text = "Low Risk"
            
            return {
                "company_name": company_name,
                "greenwashing_score": round(prob_gw, 3),
                "risk_label": risk_label,
                "model_label": risk_label, # Use simple label for UI mapping
                "details": {
                    "model_used": "Ensemble Voting Classifier",
                    "confidence": round(max(pred_proba) * 100, 1),
                    "features": features_dict # Return original features for UI
                }
            }
            
        except Exception as e:
            print(f"[ML] Ensemble prediction failed: {e}")
            # Fallback to heuristic below
            
    # 2. Heuristic Fallback (Text-based)
    sentences = re.split(r'(?<=[.!?]) +', text)
    vague_score = calculate_vague_score(sentences)
    concrete_score = calculate_concrete_score(sentences)
    sentiment = analyze_sentiment([text]) 
    
    risk_score = 0.5 + (vague_score * 0.4) - (concrete_score * 0.5)
    if sentiment['label'] == 'Negative':
        risk_score += sentiment['score'] * 0.2
    risk_score = max(0, min(1, risk_score))
    
    return {
        "company_name": company_name,
        "greenwashing_score": round(risk_score, 3),
        "risk_label": "High Risk" if risk_score > 0.7 else "Low Risk",
        "model_label": "Heuristic Analysis",
        "details": {
            "vague_language_ratio": round(vague_score, 3),
            "concrete_claims_ratio": round(concrete_score, 3),
            "model_used": "Heuristic Fallback"
        }
    }