"""
Generate mock price paths for oTree investment experiment.

This script generates pairs of graphs for n scenarios:
- graph{i}.png: Shows last 40 periods of returns and signal (from 45 generated)
- graph{i}_r.png: Shows last 41 periods (40 + one realized period)

Uses AR(1) process with parameters from the experiment:
- mu = 0.07 (mean return)
- sigma = 0.2 (volatility)
- phi = 0.7 (autocorrelation coefficient)

Generates 45 periods total, but only plots the last 40 (base) or 41 (realized).
"""

import numpy as np
import matplotlib.pyplot as plt
import os
import csv


# Parameters from the experiment
MU = 0.07  # Mean return (unconditional mean of AR(1) process)
TARGET_VOLATILITY = 0.2  # Target unconditional volatility of returns
PHI = 0.5  # Autocorrelation coefficient for AR(1)
SIGMA = TARGET_VOLATILITY * np.sqrt(1 - PHI**2)  # Shock volatility to achieve target unconditional volatility
N_PERIODS_GENERATE = 45  # Total periods to generate
N_PERIODS_PLOT = 40  # Number of periods to plot (last 40 of the 45)
N_SCENARIOS = 32  # Number of scenarios to generate


def generate_ar1_returns(n_periods, mu, sigma, phi, seed=None):
    """
    Generate AR(1) return series.
    
    Formula: r_t = mu + phi * (r_{t-1} - mu) + sigma * epsilon_t
    where epsilon_t ~ N(0, 1)
    
    Args:
        n_periods (int): Number of periods to generate
        mu (float): Mean return
        sigma (float): Volatility (standard deviation of shocks)
        phi (float): Autocorrelation coefficient
        seed (int): Random seed for reproducibility
    
    Returns:
        np.ndarray: Array of returns
    """
    if seed is not None:
        np.random.seed(seed)
    
    returns = np.zeros(n_periods)
    # Start at the mean (unconditional mean of AR(1) process)
    returns[0] = mu
    
    # Generate AR(1) process
    for t in range(1, n_periods):
        epsilon = np.random.randn()  # Standard normal shock
        returns[t] = mu + phi * (returns[t-1] - mu) + sigma * epsilon
    
    return returns


def generate_signal(returns, mu, phi):
    """
    Generate the predictive signal (one-step-ahead forecast).
    
    The signal at time t is the expected value of r_t given r_{t-1}:
    signal_t = E[r_t | r_{t-1}] = mu + phi * (r_{t-1} - mu)
    
    Args:
        returns (np.ndarray): Array of realized returns
        mu (float): Mean return
        phi (float): Autocorrelation coefficient
    
    Returns:
        np.ndarray: Array of signals (same length as returns)
    """
    signal = np.zeros(len(returns))
    
    # Signal at t=0 is just the mean (no prior information)
    signal[0] = mu
    
    # Signal at time t is the forecast made at t-1 for time t
    for t in range(1, len(returns)):
        signal[t] = mu + phi * (returns[t-1] - mu)
    
    return signal


def plot_price_path(returns, signal, graph_id, output_folder, is_realized=False, next_signal=None):
    """
    Create and save a plot of returns and signal.
    
    Args:
        returns (np.ndarray): Array of returns
        signal (np.ndarray): Array of signals
        graph_id (int): Graph identifier (1-30)
        output_folder (str): Path to output folder
        is_realized (bool): If True, save as graph{id}_r.png, else graph{id}.png
        next_signal (float): Optional. Signal prediction for next period (used in base graph)
    """
    fig, ax = plt.subplots(figsize=(8, 4.5))
    
    # Time axis (x-axis goes from right to left as in the original image)
    n_periods = len(returns)
    time_axis = np.arange(-n_periods, 0)  # E.g., -40 to -1 for 40 periods
    
    # Convert returns to percentages for plotting
    returns_pct = returns * 100
    signal_pct = signal * 100
    
    # Calculate dynamic y-axis range (2x the min-max difference)
    all_values = np.concatenate([returns_pct, signal_pct])
    if next_signal is not None:
        all_values = np.concatenate([all_values, [next_signal * 100]])
    data_min = np.min(all_values)
    data_max = np.max(all_values)
    data_range = data_max - data_min
    y_min = data_min - data_range / 2
    y_max = data_max + data_range / 2
    
    # Plot lines
    ax.plot(time_axis, returns_pct, 'r-', linewidth=2, label='Asset Return')
    ax.plot(time_axis, signal_pct, 'b--', linewidth=2, label='Signal')
    
    # Add next signal prediction in base graph
    if next_signal is not None and not is_realized:
        next_signal_pct = next_signal * 100
        ax.plot([time_axis[-1], 0], [signal_pct[-1], next_signal_pct], 'b--', linewidth=2)
        ax.scatter([0], [next_signal_pct], 
                   color='darkblue', s=60, zorder=5, marker='o', 
                   edgecolors='black', linewidths=1.5,
                   label='Signal Prediction')
    
    # Highlight last point in realized graph (smaller dot)
    if is_realized and len(time_axis) > 40:
        ax.scatter([time_axis[-1]], [returns_pct[-1]], 
                   color='darkred', s=60, zorder=5, marker='o', 
                   edgecolors='black', linewidths=1.5,
                   label='Realized Return')
    
    # Styling
    ax.set_xlabel('', fontsize=10)
    ax.set_ylabel('', fontsize=10)
    ax.set_ylim(y_min, y_max)
    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f'{y:.0f}%'))
    ax.yaxis.tick_right()
    ax.yaxis.set_label_position('right')
    
    # Add grid
    ax.grid(True, linestyle='-', alpha=0.3, color='gray')
    ax.set_axisbelow(True)
    
    # Add legend
    ax.legend(loc='upper left', fontsize=10, framealpha=0.9)
    
    # Format ticks
    ax.tick_params(axis='both', which='major', labelsize=9)
    
    # Tight layout
    plt.tight_layout()
    
    # Save figure
    if is_realized:
        filename = f"graph{graph_id}_r.png"
    else:
        filename = f"graph{graph_id}.png"
    
    filepath = os.path.join(output_folder, filename)
    plt.savefig(filepath, dpi=100, bbox_inches='tight')
    plt.close()
    
    print(f"  ✓ Saved {filename}")


def generate_scenario(scenario_id, output_folder):
    """
    Generate a complete scenario: base graph (40 periods plotted) and realized graph (41 periods plotted).
    Generates 45 periods total, then plots the last 40 for base and last 41 for realized.
    
    Args:
        scenario_id (int): Scenario identifier (1-32)
        output_folder (str): Path to output folder
    
    Returns:
        tuple: (realized_return, best_forecast) for CSV generation
    """
    print(f"Generating scenario {scenario_id}...")
    
    # Generate 46 periods total (45 + 1 realized)
    returns_full = generate_ar1_returns(
        n_periods=N_PERIODS_GENERATE + 1,
        mu=MU,
        sigma=SIGMA,
        phi=PHI,
        seed=scenario_id * 100  # Unique seed for each scenario
    )
    
    signal_full = generate_signal(returns_full, mu=MU, phi=PHI)
    
    # Base graph: last 40 periods of the 46 generated
    returns_base = returns_full[-(N_PERIODS_PLOT + 1):-1]
    signal_base = signal_full[-(N_PERIODS_PLOT + 1):-1]
    next_signal = signal_full[-1]  # Signal prediction for next period
    plot_price_path(returns_base, signal_base, scenario_id, output_folder, 
                   is_realized=False, next_signal=next_signal)
    
    # Realized graph: last 41 periods (40 base + 1 realized)
    returns_realized = returns_full[-N_PERIODS_PLOT - 1:]
    signal_realized = signal_full[-N_PERIODS_PLOT - 1:]
    plot_price_path(returns_realized, signal_realized, scenario_id, output_folder, is_realized=True)
    
    # Return data for CSV: realized return and best forecast (signal prediction)
    realized_return = returns_full[-1]
    best_forecast = signal_full[-1]
    
    return realized_return, best_forecast


def generate_all_scenarios(n_scenarios, output_folder):
    """
    Generate all scenarios and create CSV file.
    
    Args:
        n_scenarios (int): Number of scenarios to generate
        output_folder (str): Path to output folder
    """
    # Create output folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)
    
    print(f"Generating {n_scenarios} scenarios...")
    print(f"Parameters: mu={MU}, sigma={SIGMA}, phi={PHI}")
    print(f"Generate {N_PERIODS_GENERATE} periods, plot last {N_PERIODS_PLOT} (base) and last {N_PERIODS_PLOT + 1} (realized)")
    print()
    
    csv_data = []
    for i in range(1, n_scenarios + 1):
        realized_return, best_forecast = generate_scenario(i, output_folder)
        csv_data.append({
            'decision_id': i,
            'realized_return': realized_return,
            'best_forecast': best_forecast
        })
    
    # Generate CSV file
    csv_path = os.path.join(output_folder, 'realized_returns.csv')
    with open(csv_path, 'w', newline='') as csvfile:
        fieldnames = ['decision_id', 'realized_return', 'best_forecast']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(csv_data)
    
    print()
    print(f"✅ Generated {n_scenarios * 2} images ({n_scenarios} scenarios)")
    print(f"✅ Generated CSV file: realized_returns.csv")
    print(f"   Output folder: {output_folder}")


def main():
    """Main function to generate mock price paths."""
    # Get the directory of this script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    output_folder = os.path.join(script_dir, "output")
    
    # Generate all scenarios
    generate_all_scenarios(N_SCENARIOS, output_folder)


if __name__ == "__main__":
    main()