!pip install numpy pandas matplotlib seaborn --quiet

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline
plt.rcParams['figure.figsize'] = (14, 5)
plt.rcParams['axes.spines.top']   = False
plt.rcParams['axes.spines.right'] = False
print('Imports ready.')

# --- Configuration ---
TARGET_WEIGHTS = {'BTC': 0.50, 'ETH': 0.30, 'SOL': 0.20}  # must sum to 1.0
assert abs(sum(TARGET_WEIGHTS.values()) - 1.0) < 1e-9, 'Weights must sum to 1.0'

INITIAL_CAPITAL    = 10_000.0  # USD
SIMULATION_DAYS    = 365
DRIFT_THRESHOLD    = 0.05      # rebalance when any asset drifts >5% from target
REBAL_FREQUENCY    = 'ME'      # calendar rebalancing: 'W' (weekly), 'ME' (monthly)
TX_COST_BPS        = 20        # transaction cost in basis points (0.20%)
print('Config ready.')

def generate_multi_asset_prices(
    assets: List[str],
    n_days: int,
    seed: int = 42
) -> pd.DataFrame:
    """
    Generate correlated synthetic daily prices for multiple crypto assets.

    Parameters
    ----------
    assets : List[str]  Asset names (e.g. ['BTC', 'ETH', 'SOL']).
    n_days : int  Number of trading days.
    seed : int  Random seed.

    Returns
    -------
    pd.DataFrame  Daily closing prices, indexed by date, one column per asset.

    Notes
    -----
    Assets share a common market factor (BTC-like) with idiosyncratic noise.
    Altcoins have higher vol than BTC, mimicking real crypto market structure.
    """
    rng = np.random.default_rng(seed)
    dt  = 1 / 365

    # Asset parameters: (start_price, annual_vol, drift)
    params = {
        'BTC': (40_000, 0.60, 0.20),
        'ETH': (2_500,  0.75, 0.25),
        'SOL': (100,    0.90, 0.35),
        'BNB': (300,    0.70, 0.20),
        'ADA': (0.5,    0.85, 0.15),
    }

    # Shared market factor
    market_factor = rng.standard_normal(n_days)

    prices = {}
    for asset in assets:
        p0, vol, drift = params.get(asset, (100, 0.70, 0.20))
        idiosync = rng.standard_normal(n_days)
        combined = 0.7 * market_factor + 0.3 * idiosync  # 70% market, 30% idiosyncratic
        rets = (drift - 0.5 * vol**2) * dt + vol * np.sqrt(dt) * combined
        prices[asset] = p0 * np.exp(np.cumsum(rets))

    index = pd.date_range('2024-01-01', periods=n_days, freq='D')
    return pd.DataFrame(prices, index=index)


prices_df = generate_multi_asset_prices(list(TARGET_WEIGHTS.keys()), SIMULATION_DAYS)
print('Generated prices:')
print(prices_df.describe().round(0))

def simulate_portfolio(
    prices_df: pd.DataFrame,
    target_weights: Dict[str, float],
    initial_capital: float,
    rebal_strategy: str = 'threshold',
    drift_threshold: float = 0.05,
    calendar_freq: str = 'ME',
    tx_cost_bps: float = 20
) -> pd.DataFrame:
    """
    Simulate a rebalancing portfolio strategy.

    Parameters
    ----------
    prices_df : pd.DataFrame  Daily prices for each asset.
    target_weights : Dict[str, float]  Target allocation per asset.
    initial_capital : float  Starting USD capital.
    rebal_strategy : str
        'threshold', 'calendar', 'none' (drift only).
    drift_threshold : float
        Absolute weight deviation that triggers rebalancing (threshold mode).
    calendar_freq : str  Pandas offset for calendar rebalancing.
    tx_cost_bps : float  One-way transaction cost in basis points.

    Returns
    -------
    pd.DataFrame
        Daily portfolio value, weights, number of rebalancing events, total costs.
    """
    assets = list(target_weights.keys())
    tx_cost_pct = tx_cost_bps / 10_000

    # Initial allocation
    holdings = {a: (initial_capital * target_weights[a]) / prices_df[a].iloc[0] for a in assets}

    # Calendar rebalancing dates
    cal_dates = set(prices_df.resample(calendar_freq).last().index.date) if rebal_strategy == 'calendar' else set()

    records = []
    total_tx_cost = 0.0
    n_rebalances  = 0

    for date, row in prices_df.iterrows():
        # Portfolio value per asset
        values  = {a: holdings[a] * row[a] for a in assets}
        total_v = sum(values.values())
        weights = {a: values[a] / total_v for a in assets}

        # Check rebalancing trigger
        need_rebal = False
        if rebal_strategy == 'threshold':
            need_rebal = any(abs(weights[a] - target_weights[a]) > drift_threshold for a in assets)
        elif rebal_strategy == 'calendar':
            need_rebal = date.date() in cal_dates

        if need_rebal:
            n_rebalances += 1
            for a in assets:
                target_v    = total_v * target_weights[a]
                current_v   = values[a]
                trade_v     = abs(target_v - current_v)
                cost        = trade_v * tx_cost_pct
                total_tx_cost += cost
                total_v      -= cost
                holdings[a]   = (total_v * target_weights[a]) / row[a]

        records.append({
            'date':          date,
            'portfolio_value': total_v,
            'rebalanced':    need_rebal,
            'n_rebalances':  n_rebalances,
            'tx_costs_total':total_tx_cost,
            **{f'weight_{a}': weights[a] for a in assets},
        })

    return pd.DataFrame(records).set_index('date')


sim_threshold = simulate_portfolio(prices_df, TARGET_WEIGHTS, INITIAL_CAPITAL, 'threshold', DRIFT_THRESHOLD, tx_cost_bps=TX_COST_BPS)
sim_calendar  = simulate_portfolio(prices_df, TARGET_WEIGHTS, INITIAL_CAPITAL, 'calendar',  calendar_freq=REBAL_FREQUENCY, tx_cost_bps=TX_COST_BPS)
sim_none      = simulate_portfolio(prices_df, TARGET_WEIGHTS, INITIAL_CAPITAL, 'none')

for name, df in [('Threshold', sim_threshold), ('Calendar', sim_calendar), ('No Rebal', sim_none)]:
    final = df.iloc[-1]
    print(f'{name}: Value=${final["portfolio_value"]:,.0f}, Rebalances={int(final["n_rebalances"])}, TX Costs=${final["tx_costs_total"]:,.0f}')

def compute_max_drift(
    sim_df: pd.DataFrame,
    target_weights: Dict[str, float]
) -> pd.Series:
    """
    Compute the maximum absolute weight deviation across all assets per day.

    Parameters
    ----------
    sim_df : pd.DataFrame  Portfolio simulation output.
    target_weights : Dict[str, float]  Target weights.

    Returns
    -------
    pd.Series  Daily max drift from target.
    """
    drifts = pd.DataFrame()
    for a in target_weights:
        col = f'weight_{a}'
        if col in sim_df.columns:
            drifts[a] = (sim_df[col] - target_weights[a]).abs()
    return drifts.max(axis=1)


drift_threshold = compute_max_drift(sim_threshold, TARGET_WEIGHTS)
drift_calendar  = compute_max_drift(sim_calendar,  TARGET_WEIGHTS)
drift_none      = compute_max_drift(sim_none,       TARGET_WEIGHTS)

print(f'Max drift (threshold strategy): {drift_threshold.max():.1%}')
print(f'Max drift (calendar strategy):  {drift_calendar.max():.1%}')
print(f'Max drift (no rebalancing):     {drift_none.max():.1%}')

def plot_rebalancing_analysis(
    sim_threshold, sim_calendar, sim_none,
    drift_threshold, drift_calendar, drift_none
) -> None:
    """
    Three-panel plot: portfolio values, weight drift, BTC weight drift over time.

    Parameters
    ----------
    sim_threshold, sim_calendar, sim_none : pd.DataFrame  Strategy simulation results.
    drift_threshold, drift_calendar, drift_none : pd.Series  Daily max drift series.
    """
    fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)

    # Panel 1: Portfolio values
    axes[0].plot(sim_threshold.index, sim_threshold['portfolio_value'], color='green',     lw=1.5, label='Threshold')
    axes[0].plot(sim_calendar.index,  sim_calendar['portfolio_value'],  color='steelblue', lw=1.5, label='Calendar')
    axes[0].plot(sim_none.index,      sim_none['portfolio_value'],      color='gray',      lw=1.0, linestyle='--', label='No Rebalancing')
    axes[0].set_ylabel('Portfolio Value (USD)')
    axes[0].set_title('Portfolio Value: Threshold vs Calendar vs No Rebalancing')
    axes[0].legend()

    # Panel 2: Max weight drift
    axes[1].plot(drift_threshold.index, drift_threshold * 100, color='green',     lw=1.0, label='Threshold')
    axes[1].plot(drift_calendar.index,  drift_calendar  * 100, color='steelblue', lw=1.0, label='Calendar')
    axes[1].plot(drift_none.index,      drift_none      * 100, color='gray',      lw=0.8, linestyle='--', label='No Rebalancing')
    axes[1].axhline(DRIFT_THRESHOLD * 100, color='red', linewidth=0.8, linestyle=':', label=f'Threshold = {DRIFT_THRESHOLD*100:.0f}%')
    # Mark rebalancing events
    rebal_dates = sim_threshold[sim_threshold['rebalanced']].index
    axes[1].scatter(rebal_dates, [DRIFT_THRESHOLD * 100] * len(rebal_dates), color='green', s=20, zorder=5)
    axes[1].set_ylabel('Max Asset Drift (%)')
    axes[1].set_title('Maximum Weight Deviation from Target (green dots = rebalance triggered)')
    axes[1].legend()

    # Panel 3: BTC weight over time
    if 'weight_BTC' in sim_none.columns:
        axes[2].plot(sim_none.index,      sim_none['weight_BTC']      * 100, color='gray',  lw=0.8, linestyle='--', label='No Rebal')
        axes[2].plot(sim_threshold.index, sim_threshold['weight_BTC'] * 100, color='green', lw=1.5, label='Threshold')
        axes[2].axhline(TARGET_WEIGHTS['BTC'] * 100, color='red', lw=0.8, linestyle=':', label=f'Target {TARGET_WEIGHTS["BTC"]*100:.0f}%')
    axes[2].set_ylabel('BTC Weight (%)')
    axes[2].set_xlabel('Date')
    axes[2].set_title('BTC Portfolio Weight Over Time')
    axes[2].legend()

    plt.tight_layout()
    plt.show()


plot_rebalancing_analysis(sim_threshold, sim_calendar, sim_none,
                          drift_threshold, drift_calendar, drift_none)

def export_rebalancing_results(sim_threshold, sim_calendar, sim_none):
    """
    Export all rebalancing simulation results.

    Parameters
    ----------
    sim_threshold, sim_calendar, sim_none : pd.DataFrame  Strategy results.
    """
    sim_threshold.to_csv('rebal_threshold.csv')
    sim_calendar.to_csv('rebal_calendar.csv')
    sim_none.to_csv('rebal_none.csv')
    print('Exported: rebal_threshold.csv, rebal_calendar.csv, rebal_none.csv')


export_rebalancing_results(sim_threshold, sim_calendar, sim_none)