import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Any

class PlanHealthClassifier(nn.Module):
    def __init__(self, input_dim=9, hidden_dim=24):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 4)  # 4 classes: Critical, Stressed, Healthy, Excellent
        )
    def forward(self, x):
        return self.net(x)

class SentiPlanEngine:
    def __init__(self, weights_path: str):
        self.model = PlanHealthClassifier()
        if os.path.exists(weights_path):
            self.model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
        self.model.eval()

    def predict(self, text: str) -> Dict[str, Any]:
        raw = text.lower()
        
        feat = [
            len(raw) / 10000.0,
            1.0 if "income" in raw else 0.0,
            1.0 if "budget" in raw else 0.0,
            1.0 if "saving" in raw else 0.0,
            1.0 if "expense" in raw else 0.0,
            1.0 if "runway" in raw else 0.0,
            1.0 if "invest" in raw else 0.0,
            1.0 if "debt" in raw else 0.0,
            1.0 if "plan" in raw else 0.0,
            1.0 if "cashflow" in raw else 0.0
        ]
        
        with torch.no_grad():
            x = torch.tensor([feat[:9]], dtype=torch.float32)
            logits = self.model(x)
            probs = F.softmax(logits, dim=1).numpy()[0]
            pred_class = int(logits.argmax(dim=1).item())
            
        health_statuses = ["Critical / Under severe distress", "Stressed / Action needed", "Healthy / Balanced", "Excellent / High savings"]
        recommended_actions = []
        
        if pred_class == 0:
            recommended_actions.append("Reduce non-essential expenses immediately. Establish an emergency fund.")
        elif pred_class == 1:
            recommended_actions.append("Optimize budget structure. Try implementing the 50/30/20 rule to rebuild runway.")
        elif pred_class == 2:
            recommended_actions.append("Maintain current saving habits. Look into investing excess cash in money market funds.")
        elif pred_class == 3:
            recommended_actions.append("Invest surplus aggressively in diversified portfolios. Plan for long-term goals.")

        return {
            "financial_health_status": health_statuses[pred_class],
            "confidence": float(probs[pred_class]),
            "recommended_actions": recommended_actions,
            "framework": "Personal Financial Planning Standards"
        }


# ── RLM Integration ──────────────────────────────────────────────

class SentiPlanRLM:
    """
    SentiPlanRLM wraps the shared RLMEngine to provide deep reasoning capabilities
    using a dedicated Ollama specialist model (senti-plan-rlm).
    """
    def __init__(self, model_name: str = "senti-plan-rlm"):
        import sys
        import os
        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        if base_dir not in sys.path:
            sys.path.insert(0, base_dir)
            
        from senti.core.engines.superpacks.rlm_engine import RLMEngine
        self.engine = RLMEngine(model=model_name)

    async def predict_deep(self, text: str, tier: str = "C") -> dict:
        import time
        context = {
            "tier": tier,
            "domain": "sentiplan",
            "timestamp": time.time(),
        }
        
        system_suffix = (
            "Focus on personal financial planning, budgeting strategies (e.g., 50/30/20 rule), savings targets, emergency fund allocation, and cash flow forecasting."
        )
        
        rlm_response = await self.engine.reason(
            query=text,
            context=context,
            system_suffix=system_suffix
        )
        
        res = rlm_response.to_dict()
        
        # Map RLM decision to legacy SML fields for backward compatibility
        decision = res.get("decision", "healthy").lower()
        health_status = "Healthy / Balanced"
        recommended_actions = [res.get("justification", "Maintain balanced planning.")]
        
        if any(w in decision for w in ["critical", "danger", "poor", "severe", "distress"]):
            health_status = "Critical / Under severe distress"
        elif any(w in decision for w in ["stress", "warning", "action", "tight"]):
            health_status = "Stressed / Action needed"
        elif any(w in decision for w in ["excellent", "superb", "great", "surplus", "growth"]):
            health_status = "Excellent / High savings"
            
        res["financial_health_status"] = health_status
        res["confidence"] = res.get("confidence", 0.5)
        res["recommended_actions"] = recommended_actions
        res["framework"] = "Personal Financial Planning Standards"
        return res