"""
Policy Impact Simulator Service
Analyzes policy changes and their financial/social impact over time periods
"""

import json
import logging
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime, timedelta
import pandas as pd
import numpy as np
from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)

class PolicyParameter(Enum):
    """Policy parameters that can be simulated"""
    PENSION_FACTOR = "pension_factor"
    DEARNESS_RELIEF = "dearness_relief"
    BASIC_PENSION = "basic_pension"
    MINIMUM_PENSION = "minimum_pension"
    COMMUTATION_FACTOR = "commutation_factor"
    MEDICAL_ALLOWANCE = "medical_allowance"
    FAMILY_PENSION = "family_pension"
    GRATUITY_LIMIT = "gratuity_limit"
    HRA_PERCENTAGE = "hra_percentage"
    AGE_LIMIT = "retirement_age"

class ScenarioType(Enum):
    """Types of scenarios for impact analysis"""
    BEST_CASE = "best"
    BASE_CASE = "base"
    WORST_CASE = "worst"

@dataclass
class PolicyScenario:
    """Configuration for a policy scenario"""
    parameter: PolicyParameter
    current_value: float
    proposed_value: float
    effective_date: datetime
    affected_population: int
    annual_growth_rate: float = 0.03  # Default 3% annual growth
    inflation_rate: float = 0.06  # Default 6% inflation

@dataclass
class ImpactProjection:
    """Financial impact projection"""
    year: int
    baseline_cost: float
    scenario_cost: float
    impact: float
    affected_beneficiaries: int
    notes: str

class PolicyImpactSimulator:
    """Simulates financial and social impact of policy changes"""
    
    def __init__(self):
        self.policy_definitions = self._load_policy_definitions()
        self.historical_data = self._load_historical_data()
        
    def _load_policy_definitions(self) -> Dict[PolicyParameter, Dict]:
        """Load policy parameter definitions and constraints"""
        return {
            PolicyParameter.PENSION_FACTOR: {
                "name": "Pension Factor",
                "description": "Multiplier for calculating monthly pension",
                "unit": "multiplier",
                "typical_range": (0.5, 2.0),
                "current_rajasthan": 1.0,
                "impact_type": "direct_benefit"
            },
            PolicyParameter.DEARNESS_RELIEF: {
                "name": "Dearness Relief (DR)",
                "description": "Percentage increase to counter inflation",
                "unit": "percentage",
                "typical_range": (0, 50),
                "current_rajasthan": 12.0,
                "impact_type": "cost_adjustment"
            },
            PolicyParameter.BASIC_PENSION: {
                "name": "Basic Pension Amount",
                "description": "Minimum pension amount per month",
                "unit": "rupees",
                "typical_range": (3000, 15000),
                "current_rajasthan": 6000,
                "impact_type": "direct_benefit"
            },
            PolicyParameter.MINIMUM_PENSION: {
                "name": "Minimum Pension Guarantee",
                "description": "Guaranteed minimum pension amount",
                "unit": "rupees",
                "typical_range": (2000, 10000),
                "current_rajasthan": 3500,
                "impact_type": "safety_net"
            },
            PolicyParameter.MEDICAL_ALLOWANCE: {
                "name": "Medical Allowance",
                "description": "Monthly medical expense allowance",
                "unit": "rupees",
                "typical_range": (500, 5000),
                "current_rajasthan": 1000,
                "impact_type": "healthcare_benefit"
            }
        }
    
    def _load_historical_data(self) -> Dict:
        """Load historical policy implementation data"""
        return {
            "rajasthan_pensioners": {
                "2020": 450000,
                "2021": 465000,
                "2022": 480000,
                "2023": 495000,
                "2024": 510000
            },
            "average_pension": {
                "2020": 8500,
                "2021": 9200,
                "2022": 9800,
                "2023": 10400,
                "2024": 11000
            },
            "budget_allocation": {
                "2020": 3850000000,  # 385 crores
                "2021": 4280000000,  # 428 crores
                "2022": 4700000000,  # 470 crores
                "2023": 5100000000,  # 510 crores
                "2024": 5600000000   # 560 crores
            }
        }
    
    def simulate_policy_impact(
        self, 
        scenario: PolicyScenario, 
        years: int = 5,
        include_variants: bool = True
    ) -> Dict[str, Any]:
        """
        Simulate the impact of a policy change over specified years
        
        Args:
            scenario: Policy scenario configuration
            years: Number of years to project
            include_variants: Include best/worst case scenarios
        
        Returns:
            Complete impact analysis with projections and metadata
        """
        try:
            # Generate baseline projections
            baseline_projections = self._generate_baseline_projections(scenario, years)
            
            # Generate scenario projections
            scenario_projections = self._generate_scenario_projections(scenario, years)
            
            # Calculate variants if requested
            variants = {}
            if include_variants:
                variants = self._generate_scenario_variants(scenario, years)
            
            # Generate clause differences and timeline
            clause_analysis = self._analyze_clause_changes(scenario)
            
            # Create evidence pack
            evidence_pack = self._create_evidence_pack(
                scenario, baseline_projections, scenario_projections, variants
            )
            
            return {
                "scenario_id": f"policy_sim_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                "parameter": scenario.parameter.value,
                "parameter_name": self.policy_definitions[scenario.parameter]["name"],
                "current_value": scenario.current_value,
                "proposed_value": scenario.proposed_value,
                "effective_date": scenario.effective_date.isoformat(),
                "projection_years": years,
                "baseline_projections": baseline_projections,
                "scenario_projections": scenario_projections,
                "variants": variants,
                "total_impact": self._calculate_total_impact(baseline_projections, scenario_projections),
                "clause_analysis": clause_analysis,
                "evidence_pack": evidence_pack,
                "generated_at": datetime.now().isoformat(),
                "assumptions": self._document_assumptions(scenario)
            }
            
        except Exception as e:
            logger.error(f"Policy simulation error: {e}")
            return {"error": str(e)}
    
    def _generate_baseline_projections(self, scenario: PolicyScenario, years: int) -> List[ImpactProjection]:
        """Generate baseline (no change) projections"""
        projections = []
        base_population = scenario.affected_population
        current_avg_benefit = self._estimate_current_benefit(scenario)
        
        for year in range(1, years + 1):
            # Account for population growth and inflation
            year_population = int(base_population * (1 + scenario.annual_growth_rate) ** year)
            year_benefit = current_avg_benefit * (1 + scenario.inflation_rate) ** year
            
            annual_cost = year_population * year_benefit * 12  # Monthly to annual
            
            projections.append(ImpactProjection(
                year=year,
                baseline_cost=annual_cost,
                scenario_cost=annual_cost,  # Same for baseline
                impact=0,
                affected_beneficiaries=year_population,
                notes=f"Baseline year {year}: No policy change, inflation adjustment only"
            ))
        
        return projections
    
    def _generate_scenario_projections(self, scenario: PolicyScenario, years: int) -> List[ImpactProjection]:
        """Generate scenario (with change) projections"""
        projections = []
        base_population = scenario.affected_population
        current_benefit = self._estimate_current_benefit(scenario)
        new_benefit = self._calculate_new_benefit(scenario)
        
        for year in range(1, years + 1):
            year_population = int(base_population * (1 + scenario.annual_growth_rate) ** year)
            
            # Baseline cost (what it would have been)
            baseline_benefit = current_benefit * (1 + scenario.inflation_rate) ** year
            baseline_cost = year_population * baseline_benefit * 12
            
            # Scenario cost (with policy change)
            scenario_benefit = new_benefit * (1 + scenario.inflation_rate) ** year
            scenario_cost = year_population * scenario_benefit * 12
            
            impact = scenario_cost - baseline_cost
            
            projections.append(ImpactProjection(
                year=year,
                baseline_cost=baseline_cost,
                scenario_cost=scenario_cost,
                impact=impact,
                affected_beneficiaries=year_population,
                notes=f"Year {year}: Policy change impact ₹{impact/10000000:.1f} crores"
            ))
        
        return projections
    
    def _generate_scenario_variants(self, scenario: PolicyScenario, years: int) -> Dict[str, List[ImpactProjection]]:
        """Generate best/worst case scenario variants"""
        variants = {}
        
        # Best case: Lower growth, higher efficiency
        best_scenario = PolicyScenario(
            parameter=scenario.parameter,
            current_value=scenario.current_value,
            proposed_value=scenario.proposed_value,
            effective_date=scenario.effective_date,
            affected_population=int(scenario.affected_population * 0.9),  # 10% fewer beneficiaries
            annual_growth_rate=scenario.annual_growth_rate * 0.8,  # 20% lower growth
            inflation_rate=scenario.inflation_rate * 0.9  # 10% lower inflation
        )
        variants["best_case"] = self._generate_scenario_projections(best_scenario, years)
        
        # Worst case: Higher growth, implementation challenges
        worst_scenario = PolicyScenario(
            parameter=scenario.parameter,
            current_value=scenario.current_value,
            proposed_value=scenario.proposed_value * 1.1,  # 10% higher due to implementation costs
            effective_date=scenario.effective_date,
            affected_population=int(scenario.affected_population * 1.2),  # 20% more beneficiaries
            annual_growth_rate=scenario.annual_growth_rate * 1.3,  # 30% higher growth
            inflation_rate=scenario.inflation_rate * 1.1  # 10% higher inflation
        )
        variants["worst_case"] = self._generate_scenario_projections(worst_scenario, years)
        
        return variants
    
    def _analyze_clause_changes(self, scenario: PolicyScenario) -> Dict[str, Any]:
        """Analyze what policy clauses would change"""
        parameter_def = self.policy_definitions[scenario.parameter]
        
        return {
            "affected_clauses": self._identify_affected_clauses(scenario.parameter),
            "clause_diff": {
                "before": f"{parameter_def['name']}: {scenario.current_value} {parameter_def['unit']}",
                "after": f"{parameter_def['name']}: {scenario.proposed_value} {parameter_def['unit']}",
                "change_type": "increase" if scenario.proposed_value > scenario.current_value else "decrease",
                "change_magnitude": abs(scenario.proposed_value - scenario.current_value),
                "change_percentage": ((scenario.proposed_value - scenario.current_value) / scenario.current_value) * 100
            },
            "effective_timeline": {
                "announcement_date": (scenario.effective_date - timedelta(days=90)).isoformat(),
                "legislative_process": (scenario.effective_date - timedelta(days=60)).isoformat(),
                "notification_date": (scenario.effective_date - timedelta(days=30)).isoformat(),
                "effective_date": scenario.effective_date.isoformat(),
                "first_payment": (scenario.effective_date + timedelta(days=30)).isoformat()
            },
            "legal_references": self._get_legal_references(scenario.parameter)
        }
    
    def _identify_affected_clauses(self, parameter: PolicyParameter) -> List[str]:
        """Identify which policy clauses would be affected"""
        clause_mapping = {
            PolicyParameter.PENSION_FACTOR: [
                "Rajasthan Civil Services (Pension) Rules, 1996 - Rule 35",
                "Pension calculation formula - Clause 4.2.1",
                "Service weightage provisions - Clause 6.1"
            ],
            PolicyParameter.DEARNESS_RELIEF: [
                "Dearness Relief calculation - Rule 42",
                "Inflation adjustment mechanism - Clause 3.4",
                "Automatic revision provisions - Rule 43"
            ],
            PolicyParameter.BASIC_PENSION: [
                "Minimum pension guarantee - Rule 28",
                "Basic pension structure - Clause 2.1",
                "Pension floor provisions - Rule 29"
            ],
            PolicyParameter.MEDICAL_ALLOWANCE: [
                "Medical benefits - Rule 52",
                "Healthcare allowance - Clause 8.3",
                "Medical reimbursement - Rule 53"
            ]
        }
        return clause_mapping.get(parameter, ["General pension provisions"])
    
    def _get_legal_references(self, parameter: PolicyParameter) -> List[str]:
        """Get relevant legal references for the parameter"""
        references = {
            PolicyParameter.PENSION_FACTOR: [
                "Rajasthan Civil Services (Pension) Rules, 1996",
                "Rajasthan Government Resolution No. F.2(5)FD/Rules/96",
                "Central Civil Services (Pension) Rules, 2021 - Reference"
            ],
            PolicyParameter.DEARNESS_RELIEF: [
                "Rajasthan Civil Services (Revised Pay) Rules, 2017",
                "Dearness Allowance calculation guidelines",
                "RCS(RP) Rules notification dated 15.08.2017"
            ]
        }
        return references.get(parameter, ["Rajasthan Civil Services (Pension) Rules, 1996"])
    
    def _estimate_current_benefit(self, scenario: PolicyScenario) -> float:
        """Estimate current monthly benefit amount"""
        parameter_def = self.policy_definitions[scenario.parameter]
        
        if scenario.parameter == PolicyParameter.BASIC_PENSION:
            return scenario.current_value
        elif scenario.parameter == PolicyParameter.PENSION_FACTOR:
            # Average salary * factor
            return 25000 * scenario.current_value  # Assumed average salary
        elif scenario.parameter == PolicyParameter.DEARNESS_RELIEF:
            # DR percentage of basic pension
            base_pension = 8000  # Assumed base
            return base_pension * (scenario.current_value / 100)
        else:
            return parameter_def.get("current_rajasthan", 5000)
    
    def _calculate_new_benefit(self, scenario: PolicyScenario) -> float:
        """Calculate new monthly benefit amount after policy change"""
        if scenario.parameter == PolicyParameter.BASIC_PENSION:
            return scenario.proposed_value
        elif scenario.parameter == PolicyParameter.PENSION_FACTOR:
            return 25000 * scenario.proposed_value
        elif scenario.parameter == PolicyParameter.DEARNESS_RELIEF:
            base_pension = 8000
            return base_pension * (scenario.proposed_value / 100)
        else:
            return scenario.proposed_value
    
    def _calculate_total_impact(self, baseline: List[ImpactProjection], scenario: List[ImpactProjection]) -> Dict[str, float]:
        """Calculate total financial impact"""
        total_additional_cost = sum(proj.impact for proj in scenario)
        total_baseline_cost = sum(proj.baseline_cost for proj in baseline)
        
        return {
            "total_additional_cost_crores": total_additional_cost / 10000000,  # Convert to crores
            "total_baseline_cost_crores": total_baseline_cost / 10000000,
            "percentage_increase": (total_additional_cost / total_baseline_cost) * 100 if total_baseline_cost > 0 else 0,
            "annual_average_impact_crores": (total_additional_cost / len(scenario)) / 10000000,
            "cost_per_beneficiary_annual": total_additional_cost / (scenario[0].affected_beneficiaries * len(scenario))
        }
    
    def _create_evidence_pack(self, scenario, baseline, projections, variants) -> Dict[str, Any]:
        """Create exportable evidence pack"""
        return {
            "summary_stats": {
                "policy_parameter": self.policy_definitions[scenario.parameter]["name"],
                "change_magnitude": f"{scenario.current_value} → {scenario.proposed_value}",
                "affected_population": scenario.affected_population,
                "projection_period": f"{len(projections)} years",
                "total_impact": f"₹{sum(p.impact for p in projections)/10000000:.1f} crores"
            },
            "yearly_breakdown": [
                {
                    "year": p.year,
                    "baseline_crores": p.baseline_cost / 10000000,
                    "scenario_crores": p.scenario_cost / 10000000,
                    "impact_crores": p.impact / 10000000,
                    "beneficiaries": p.affected_beneficiaries
                }
                for p in projections
            ],
            "scenario_comparison": {
                "best_case_total": sum(p.impact for p in variants.get("best_case", [])) / 10000000 if variants else 0,
                "base_case_total": sum(p.impact for p in projections) / 10000000,
                "worst_case_total": sum(p.impact for p in variants.get("worst_case", [])) / 10000000 if variants else 0
            },
            "export_formats": {
                "csv_data": "Ready for CSV export",
                "chart_data": "Ready for visualization",
                "pdf_report": "Formatted for official reporting"
            }
        }
    
    def _document_assumptions(self, scenario: PolicyScenario) -> Dict[str, Any]:
        """Document all assumptions used in the simulation"""
        return {
            "demographic_assumptions": {
                "affected_population": scenario.affected_population,
                "annual_growth_rate": f"{scenario.annual_growth_rate*100:.1f}%",
                "population_source": "Rajasthan pension department estimates"
            },
            "economic_assumptions": {
                "inflation_rate": f"{scenario.inflation_rate*100:.1f}%",
                "salary_growth": "Aligned with inflation",
                "implementation_efficiency": "100% (no leakage assumed)"
            },
            "policy_assumptions": {
                "effective_date": scenario.effective_date.strftime("%Y-%m-%d"),
                "implementation_delay": "None assumed",
                "administrative_costs": "Not included in projections",
                "compliance_rate": "100% (full implementation assumed)"
            },
            "data_sources": [
                "Rajasthan Finance Department budget documents",
                "Pension disbursement historical data",
                "National Sample Survey Office reports",
                "Reserve Bank of India inflation projections"
            ]
        }

# Usage example and helper functions
def create_sample_scenarios() -> List[PolicyScenario]:
    """Create sample policy scenarios for testing"""
    return [
        PolicyScenario(
            parameter=PolicyParameter.DEARNESS_RELIEF,
            current_value=12.0,
            proposed_value=18.0,
            effective_date=datetime(2025, 4, 1),
            affected_population=510000,
            annual_growth_rate=0.03,
            inflation_rate=0.06
        ),
        PolicyScenario(
            parameter=PolicyParameter.BASIC_PENSION,
            current_value=6000,
            proposed_value=8000,
            effective_date=datetime(2025, 7, 1),
            affected_population=450000,
            annual_growth_rate=0.025,
            inflation_rate=0.055
        ),
        PolicyScenario(
            parameter=PolicyParameter.MEDICAL_ALLOWANCE,
            current_value=1000,
            proposed_value=2000,
            effective_date=datetime(2025, 10, 1),
            affected_population=510000,
            annual_growth_rate=0.035,
            inflation_rate=0.065
        )
    ]