"""
Monte Carlo Simulation Template for Financial Modeling
=======================================================
A reusable class for any Monte Carlo simulation.
Adapt the `calculate_target()` method for your specific use case.

Usage:
    sim = MonteCarloSimulator(n_simulations=10000, seed=42)
    sim.add_variable('revenue_growth', 'triangular', 0.06, 0.10, 0.14)
    sim.add_variable('wacc', 'normal', 0.11, 0.008)
    sim.run(calculate_dcf)          # Pass your calculation function
    sim.analyze(target_value=3800)  # Optional: compare to a target
    sim.plot(title='DCF Valuation', xlabel='Value per Share (₹)')
    sim.export('results.csv')
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Callable, Optional, Dict, Any


class MonteCarloSimulator:
    """Reusable Monte Carlo simulation engine for financial modeling."""
    
    def __init__(self, n_simulations: int = 10000, seed: int = 42):
        self.n_simulations = n_simulations
        self.seed = seed
        self.variables: Dict[str, np.ndarray] = {}
        self.results: Optional[np.ndarray] = None
        self._variable_configs = {}  # Store for reporting
    
    def add_variable(self, name: str, distribution: str, *params):
        """
        Add a stochastic variable.
        
        Distributions:
            'normal'    → add_variable('x', 'normal', mean, std)
            'triangular'→ add_variable('x', 'triangular', min, mode, max)
            'uniform'   → add_variable('x', 'uniform', low, high)
            'lognormal' → add_variable('x', 'lognormal', mean, sigma)
            'fixed'     → add_variable('x', 'fixed', value)
        """
        np.random.seed(self.seed)
        
        if distribution == 'normal':
            mean, std = params
            values = np.random.normal(mean, std, self.n_simulations)
            self._variable_configs[name] = f'Normal(μ={mean}, σ={std})'
        elif distribution == 'triangular':
            left, mode, right = params
            values = np.random.triangular(left, mode, right, self.n_simulations)
            self._variable_configs[name] = f'Tri({left}, {mode}, {right})'
        elif distribution == 'uniform':
            low, high = params
            values = np.random.uniform(low, high, self.n_simulations)
            self._variable_configs[name] = f'Uniform({low}, {high})'
        elif distribution == 'lognormal':
            mean, sigma = params
            values = np.random.lognormal(mean, sigma, self.n_simulations)
            self._variable_configs[name] = f'LogN(μ={mean}, σ={sigma})'
        elif distribution == 'fixed':
            values = np.full(self.n_simulations, params[0])
            self._variable_configs[name] = f'Fixed({params[0]})'
        else:
            raise ValueError(f"Unknown distribution: {distribution}")
        
        self.variables[name] = values
    
    def run(self, calculation_fn: Callable[[Dict[str, np.ndarray], int], float]):
        """
        Run the simulation using the provided calculation function.
        
        Args:
            calculation_fn: A function that takes (variables_dict, iteration_index)
                           and returns a single numeric result.
        
        Example calculation function for DCF:
            def calc_dcf(vars, i):
                growth = vars['revenue_growth'][i]
                wacc = vars['wacc'][i]
                ...  # Your DCF logic here
                return intrinsic_value
        """
        self.results = np.zeros(self.n_simulations)
        for i in range(self.n_simulations):
            self.results[i] = calculation_fn(self.variables, i)
        return self
    
    def analyze(self, target_value: float = None) -> pd.DataFrame:
        """Print and return summary statistics."""
        if self.results is None:
            raise RuntimeError("Run the simulation first!")
        
        r = self.results
        stats = {
            'N Simulations': len(r),
            'Mean': r.mean(),
            'Median': np.median(r),
            'Std Dev': r.std(),
            'Min': r.min(),
            'Max': r.max(),
            'P5 (Pessimistic)': np.percentile(r, 5),
            'P25': np.percentile(r, 25),
            'P50 (Median)': np.percentile(r, 50),
            'P75': np.percentile(r, 75),
            'P95 (Optimistic)': np.percentile(r, 95),
        }
        
        print("=" * 55)
        print("MONTE CARLO SIMULATION RESULTS")
        print("=" * 55)
        print(f"\n{'Variable':<25} {'Distribution':<30}")
        print("-" * 55)
        for name, config in self._variable_configs.items():
            print(f"  {name:<25} {config:<30}")
        
        print(f"\n{'Statistic':<25} {'Value':>20}")
        print("-" * 55)
        for k, v in stats.items():
            if isinstance(v, float):
                print(f"  {k:<25} {v:>20,.2f}")
            else:
                print(f"  {k:<25} {v:>20}")
        
        if target_value is not None:
            prob_above = (r > target_value).mean() * 100
            print(f"\n  Target Value:            {target_value:>20,.2f}")
            print(f"  P(Result > Target):      {prob_above:>19.1f}%")
            print(f"  P(Result < Target):      {100-prob_above:>19.1f}%")
        
        return pd.DataFrame([stats])
    
    def plot(self, title: str = 'Monte Carlo Simulation Results',
             xlabel: str = 'Value', target_value: float = None,
             figsize: tuple = (14, 6), save_path: str = None):
        """Create histogram + CDF plots."""
        if self.results is None:
            raise RuntimeError("Run the simulation first!")
        
        fig, axes = plt.subplots(1, 2, figsize=figsize)
        r = self.results
        
        # Histogram
        axes[0].hist(r, bins=50, color='#3B82F6', alpha=0.7, edgecolor='white')
        axes[0].axvline(np.median(r), color='green', linewidth=2, linestyle='-',
                        label=f'Median: {np.median(r):,.0f}')
        if target_value is not None:
            color = 'red' if target_value < np.median(r) else 'orange'
            axes[0].axvline(target_value, color=color, linewidth=2, linestyle='--',
                            label=f'Target: {target_value:,.0f}')
        axes[0].set_title(f'{title} — Histogram', fontweight='bold')
        axes[0].set_xlabel(xlabel)
        axes[0].set_ylabel('Frequency')
        axes[0].legend()
        axes[0].grid(alpha=0.3)
        
        # CDF
        sorted_r = np.sort(r)
        cdf = np.arange(1, len(sorted_r) + 1) / len(sorted_r)
        axes[1].plot(sorted_r, cdf, color='#7C3AED', linewidth=2)
        if target_value is not None:
            axes[1].axvline(target_value, color='red', linewidth=2, linestyle='--',
                            label=f'Target: {target_value:,.0f}')
        axes[1].set_title(f'{title} — CDF', fontweight='bold')
        axes[1].set_xlabel(xlabel)
        axes[1].set_ylabel('Cumulative Probability')
        axes[1].legend()
        axes[1].grid(alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.show()
    
    def export(self, filepath: str = 'monte_carlo_results.csv'):
        """Export all simulation results and variables to CSV."""
        if self.results is None:
            raise RuntimeError("Run the simulation first!")
        
        df = pd.DataFrame(self.variables)
        df['Result'] = self.results
        df.to_csv(filepath, index=False)
        print(f"✅ Results exported to {filepath} ({len(df)} rows)")
        return df
    
    def tornado(self, calculation_fn, title='Tornado Chart', figsize=(10, 6)):
        """Sensitivity analysis: vary each input ±1 std dev, hold others at mean."""
        if not self.variables:
            raise RuntimeError("Add variables first!")
        
        base_vars = {k: v.mean() for k, v in self.variables.items()}
        
        sensitivities = []
        for name, values in self.variables.items():
            std = values.std()
            if std == 0:
                continue
            
            # Low: set this var to mean-1std, others at mean
            low_vars = {k: (v.mean() if k != name else v.mean() - std) 
                       for k, v in self.variables.items()}
            low_result = calculation_fn(low_vars, 0)
            
            # High: set this var to mean+1std
            high_vars = {k: (v.mean() if k != name else v.mean() + std) 
                        for k, v in self.variables.items()}
            high_result = calculation_fn(high_vars, 0)
            
            sensitivities.append({
                'Variable': name,
                'Low Result': low_result,
                'High Result': high_result,
                'Swing': abs(high_result - low_result)
            })
        
        # Sort by swing
        df_tornado = pd.DataFrame(sensitivities).sort_values('Swing')
        
        fig, ax = plt.subplots(figsize=figsize)
        for _, row in df_tornado.iterrows():
            ax.barh(row['Variable'], row['High Result'] - row['Low Result'],
                    left=row['Low Result'], color='#3B82F6', alpha=0.7, edgecolor='white')
        ax.set_title(title, fontweight='bold', fontsize=13)
        ax.set_xlabel('Output Value')
        ax.grid(alpha=0.3, axis='x')
        plt.tight_layout()
        plt.show()
        
        return df_tornado


# ========================================================
# EXAMPLE USAGE: DCF Valuation with the Template
# ========================================================
if __name__ == '__main__':
    
    # 1. Create simulator
    sim = MonteCarloSimulator(n_simulations=10000, seed=42)
    
    # 2. Define stochastic variables
    sim.add_variable('rev_growth', 'triangular', 0.06, 0.10, 0.14)
    sim.add_variable('ebitda_margin', 'normal', 0.26, 0.015)
    sim.add_variable('wacc', 'normal', 0.11, 0.008)
    sim.add_variable('terminal_growth', 'triangular', 0.04, 0.055, 0.07)
    
    # Fixed values (still added for the calculation function to use)
    sim.add_variable('base_revenue', 'fixed', 240_000)
    sim.add_variable('net_debt', 'fixed', 8_000)
    sim.add_variable('shares', 'fixed', 361)
    
    # 3. Define the DCF calculation function
    def calculate_dcf(vars, i):
        revenue = vars['base_revenue'][i]
        fcff_list = []
        for yr in range(1, 6):
            revenue *= (1 + vars['rev_growth'][i])
            ebitda = revenue * vars['ebitda_margin'][i]
            nopat = ebitda * (1 - 0.2517)
            fcff = nopat + revenue * 0.03 - revenue * 0.04 - revenue * 0.02
            fcff_list.append(fcff)
        
        terminal_value = fcff_list[-1] * (1 + vars['terminal_growth'][i]) / \
                        (vars['wacc'][i] - vars['terminal_growth'][i])
        
        pv_fcff = sum(f / (1 + vars['wacc'][i])**t 
                     for t, f in enumerate(fcff_list, 1))
        pv_tv = terminal_value / (1 + vars['wacc'][i])**5
        
        ev = pv_fcff + pv_tv
        return (ev - vars['net_debt'][i]) / vars['shares'][i]
    
    # 4. Run simulation
    sim.run(calculate_dcf)
    
    # 5. Analyze
    sim.analyze(target_value=3800)
    
    # 6. Visualize
    sim.plot(title='TCS DCF Valuation', xlabel='Intrinsic Value per Share (₹)',
             target_value=3800, save_path='template_dcf_mc.png')
    
    # 7. Export
    sim.export('tcs_monte_carlo_results.csv')
    
    # 8. Tornado chart (sensitivity)
    # sim.tornado(calculate_dcf, title='DCF Sensitivity — TCS')