#!/usr/bin/env python3
"""Monthly full-wealth share simulation with concentration/income feedback.

No 10% or 20% BTC boundary. Signals and market shocks feed a monotone implicit
price-adjustment step. Public outputs include first-passage probabilities and
sensitivity to the economic feedback, not just its central forecast.
"""
import argparse
import calendar
import csv
from datetime import datetime
import hashlib
import json
from pathlib import Path

import numpy as np
from scipy.special import expit, logit
from permanent_loss import PermanentLoss
from portfolio_models import Macro, Demand, Equilibrium
from share_feedback import economic_state, implicit_step, diffusion_clock, diffusion_scale
from statistical_trend import log_factor as curve_log_factor

ROOT=Path(__file__).resolve().parent
PROBS=np.array([.025,.05,.1,.25,.5,.75,.9,.95,.975,.99])
KEYS=[f'p{int(p*1000):03d}' for p in PROBS]
BASES=['nominal_usd','real_2026sep_usd']
THRESHOLDS=[(.01,'1'),(.05,'5'),(.075,'7_5'),(.09,'9'),(.10,'10'),(.15,'15'),(.20,'20'),(.30,'30')]


def config():return json.loads((ROOT/'path_assumptions.json').read_text())


def supply(date):
    anchor=datetime(2024,4,20,0,9,27)
    n=840001+int((date-anchor).total_seconds()//600)
    return sum(max(0,min(210000,n-k*210000))*((50*100_000_000)>>k) for k in range(64))/1e8


def monthly_dates(settings):
    anchor=datetime.fromisoformat(settings['as_of']);dates=[anchor]
    year,month=anchor.year,anchor.month
    while year<=settings['simulation_end_year']:
        date=datetime(year,month,calendar.monthrange(year,month)[1])
        if date>anchor:dates.append(date)
        year,month=(year+1,1) if month==12 else (year,month+1)
    dates.extend(datetime(year,1,1) for year in range(anchor.year+1,settings['simulation_end_year']+1))
    return sorted(set(dates))


def write_csv(directory,name,rows):
    if not rows:return
    with (directory/name).open('w',newline='') as f:
        writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)


def pair_standard_error(values):
    half=len(values)//2;pairs=(values[:half]+values[half:])/2
    return float(pairs.std(ddof=1)/np.sqrt(half))


def summarize(price,date,cpi,loss,uniform,case):
    alive=uniform>=loss;observed_loss=float(1-alive.mean());values=np.where(alive,price,0.)
    qs=np.quantile(values,PROBS,method='inverted_cdf');median=float(np.quantile(price[alive],.5,method='inverted_cdf'))
    mean=float(values.mean());conditional=float(price[alive].mean());counterfactual=float(price.mean())
    standard_error=pair_standard_error(values);positive_median=float(np.median(price));rows=[];coverage=[]
    for basis,deflator in zip(BASES,[1.,cpi]):
        rows.append(dict(case=case['case'],label_ja=case['label_ja'],date=date.isoformat(),year=date.year,basis=basis,
            scope='all_states',simulation_paths=len(price),loss_probability=loss,simulated_loss_probability=observed_loss,
            mean=mean/deflator,mean_standard_error=standard_error/deflator,conditional_mean=conditional/deflator,
            conditional_median=median/deflator,counterfactual_mean=counterfactual/deflator,counterfactual_median=positive_median/deflator,
            **dict(zip(KEYS,(qs/deflator).tolist()))))
        for low,high in [(.025,.975),(.1,.9),(.25,.75)]:
            lower,upper=[qs[np.flatnonzero(np.isclose(PROBS,p))[0]] for p in [low,high]]
            coverage.append(dict(case=case['case'],date=date.isoformat(),year=date.year,basis=basis,
                interval=f'P{100*low:g}–P{100*high:g}',actual_probability=float(((values>=lower)&(values<=upper)).mean())))
    return rows,coverage



def cases():
    c=config();base=dict(case='reference',label_ja='基準（集中・所得のフィードバック）',process='reference',
        weights=c['source_weights'],loss_model='integrated',loss_scenario='central',hazard=0.,
        diffusion_scale=1.,diffusion_months=c['diffusion_months'],trend_months=120,maturing_market_risk=True,feedback={})
    def v(key,label,same=False,**kw):
        return dict(base,**dict(kw,case=key,label_ja=label,process='reference' if same else key))
    return [base,
        v('no_loss','回復不能な価値喪失なし',True,loss_model='constant'),
        v('loss_low','回復不能な価値喪失・低仮定',True,loss_scenario='low'),
        v('loss_high','回復不能な価値喪失・高仮定',True,loss_scenario='high'),
        v('fixed_hazard','一定危険率年0.5%',True,loss_model='constant',hazard=.005),
        *[v('loss_'+m,label+'のみ',True,loss_model=m) for m,label in [('reference_class','参照クラス法'),('fault_tree','フォールトツリー法'),('hazard','時変ハザード法'),('bayesian','ベイズ法')]],
        v('statistical_only','統計のみ',weights=[1.,0.,0.]),
        v('demand_only','需要のみ',weights=[0.,1.,0.]),
        v('equilibrium_only','均衡のみ',weights=[0.,0.,1.]),
        v('stat_low','統計の係数20%',weights=[.2,.4,.4]),
        v('stat_high','統計の係数50%',weights=[.5,.25,.25]),
        v('vol_low','市場変動率0.75倍',diffusion_scale=.75),
        v('vol_high','市場変動率1.25倍',diffusion_scale=1.25),
        v('history60','統計の回帰を60か月',trend_months=60),
        v('power_age_trend','鈍化の形を固定した旧回帰',trend_family='log_age'),
        v('sigma60','変動率の推定を60か月',diffusion_months=60),
        v('response_low','価格反応の強さ0.5',feedback={'response':.5}),
        v('response_high','価格反応の強さ4',feedback={'response':4.}),
        v('risk_low','非線形費用のリスク回避度2',feedback={'risk_aversion':2.}),
        v('risk_high','非線形費用のリスク回避度6',feedback={'risk_aversion':6.}),
        v('constant_market_risk','市場変動率を起点に固定',maturing_market_risk=False),
        v('income_zero','競合所得の効果なし',feedback={'competing_income_yield':0.}),
        v('wealth_high','運用対象資産の割合60%',feedback={'investable_fraction':.60}),
        v('wealth_low','運用対象資産の割合40%',feedback={'investable_fraction':.40}),
        v('no_feedback','集中・所得フィードバックなし',feedback={'response':0.})]


def run(n,seed,directory,cache_dir=None,targets_only=False,reference_only=False,substeps=None,state_recorder=None):
    if n<1024 or n%2:raise ValueError('Use an even number of at least 1024 paths')
    settings=config();inputs=json.loads((ROOT/'portfolio_inputs.json').read_text())
    calibration=json.loads((ROOT/'monetary_inputs.json').read_text())
    numerical=dict(np.load(ROOT/'portfolio_inputs.npz'));statistical=dict(np.load(ROOT/'monetary_inputs.npz'))
    dates=monthly_dates(settings);anchor=dates[0];half=n//2
    times=np.array([(d-anchor).days/365.25 for d in dates])
    substeps=substeps or settings['integration_substeps_per_month']
    macro=Macro(settings,inputs);loss=PermanentLoss();all_cases=[c for c in cases() if c['process']=='reference'] if reference_only else cases();base=all_cases[0]
    processes=[c for c in all_cases if c['case']==c['process']];index={c['case']:k for k,c in enumerate(processes)}
    fundamental,market_rng,death_rng=[np.random.default_rng(s) for s in np.random.SeedSequence(seed).spawn(3)]
    uniform=(np.arange(n)+death_rng.random(n))/n;death_rng.shuffle(uniform)
    s0=settings['spot_usd']*supply(anchor)/macro.wealth0
    rank=fundamental.random(half)
    betas={m:np.interp(rank,np.linspace(0,1,len(statistical['beta'+str(m)])),statistical['beta'+str(m)]) for m in [120,60]}
    curve_draws={m:statistical['curve'+str(m)][np.minimum((rank*len(statistical['curve'+str(m)])).astype(int),len(statistical['curve'+str(m)])-1)] for m in [120,60]} if 'curve120' in statistical else None
    sigma=np.array([calibration['trends'][str(c['trend_months'])]['innovation_estimates'][str(c['diffusion_months'])]['annual_log_odds_diffusion']*c['diffusion_scale'] for c in processes])
    demand=Demand(half,fundamental,inputs,settings['spot_usd'],supply(anchor))
    equilibrium=Equilibrium(settings,inputs,numerical,macro,supply,dates)
    dependency={k:settings[k] for k in ['simulation_end_year','equilibrium_terminal_year','equilibrium_max_step_years','method_version']}
    digest=hashlib.sha256((ROOT/'portfolio_models.py').read_bytes()+(ROOT/'portfolio_inputs.npz').read_bytes()+
        (ROOT/'portfolio_inputs.json').read_bytes()+(ROOT/'permanent_loss_inputs.json').read_bytes()+b''.join((ROOT/n).read_bytes() for n in ['macro_path.py','asset_model.py','asset_risk_inputs.json','asset_return_inputs.json'])+json.dumps(dependency,sort_keys=True).encode()).hexdigest()[:20]
    cache=None if cache_dir is None else cache_dir/('equilibrium_'+digest+'.npy')
    print(f'Preparing {len(dates)} native monthly states ...',flush=True)
    if cache is not None and cache.exists():eq_factors=np.load(cache)
    else:
        eq_factors=equilibrium.solve(dates)
        if cache is not None:cache.parent.mkdir(parents=True,exist_ok=True);np.save(cache,eq_factors)
    eq_index=fundamental.integers(0,eq_factors.shape[1],size=half)
    age=(anchor-datetime(2009,1,3)).days/365.25
    x=np.full((len(processes),n),logit(s0));previous_f=np.zeros((len(processes),half));brownian=np.zeros(n)
    previous_alive=np.ones(n,dtype=bool)
    ever={p:np.zeros(n,dtype=bool) for p,_ in THRESHOLDS}
    ever_counterfactual={p:np.zeros(n,dtype=bool) for p,_ in THRESHOLDS}
    monthly=[];targets=[];sensitivity=[];coverage=[];probabilities=[];components=[];residual=[]
    diagnostics=[];central=[];loss_curves=[];moment_rows=[];sample_paths=[]
    max_partition_error=max_anchor_error=max_loss_error=max_root_residual=0.
    target_set={datetime(y,1,1) for y in settings['target_years']}
    print(f'Simulating {n:,} paths, {len(processes)} price processes, {substeps} steps/month ...',flush=True)
    for j,(date,t) in enumerate(zip(dates,times)):
        trend={m:(curve_log_factor(t,age,calibration['trends'][str(m)]['active_curve']['reference_age'],curve_draws[m])
                  if settings.get('statistical_family')=='power_transformed_age' else betas[m]*np.log1p(t/age)) for m in [120,60]}
        factors=np.column_stack([trend[120],demand.log_factor(t,supply(date),concentration_feedback=True),eq_factors[j,eq_index]])
        f=np.array([factors@np.asarray(c['weights']) for c in processes])
        for k,c in enumerate(processes):
            selected=(betas[c['trend_months']]*np.log1p(t/age) if c.get('trend_family')=='log_age' else trend[c['trend_months']])
            f[k]+=c['weights'][0]*(selected-trend[120])
        if not j:
            if np.max(np.abs(f))>1e-5:raise AssertionError('Native signals must start at zero')
            f[:]=0.
        else:
            dt=(t-times[j-1])/substeps
            delta=np.tile((f-previous_f)/substeps,(1,2))
            innovation_grid=8
            if innovation_grid%substeps:
                raise ValueError('Integration substeps must divide the common eight-increment monthly grid')
            normal=market_rng.standard_normal((innovation_grid,half))
            fine_times=np.linspace(times[j-1],t,innovation_grid+1)
            fine=normal*np.sqrt(np.diff(diffusion_clock(fine_times,inputs)))[:,None]
            constant_fine=normal*np.sqrt((t-times[j-1])/innovation_grid)
            for step in range(substeps):
                now=times[j-1]+(step+1)*dt
                # A fixed fine innovation grid couples substep refinements.
                width=innovation_grid//substeps
                z=fine[step*width:(step+1)*width].sum(0);innovation=np.r_[z,-z];brownian+=innovation
                hazard=loss.intensity(anchor+(date-anchor)*(now/t)) if t else loss.intensity(anchor)
                for k,c in enumerate(processes):
                    state=economic_state(now,settings,c['feedback'])
                    shock=innovation
                    if not c['maturing_market_risk']:
                        raw=constant_fine[step*width:(step+1)*width].sum(0);shock=np.r_[raw,-raw]
                    x[k],error=implicit_step(x[k],delta[k],sigma[k]*shock,dt,s0,state,hazard)
                    max_root_residual=max(max_root_residual,error)
                    if state['response'] and np.max(expit(x[k]))>=state['available']:
                        raise AssertionError('A state exhausted participating investable wealth')
        previous_f=f
        shares=expit(x);share=shares[0]
        wealth,cpi=macro.values(t);quantity=supply(date);price=wealth/quantity*share
        q=loss.probability(date);alive=uniform>=q
        if state_recorder is not None:
            # Observation hook: no extra draws or changes to the price process.
            state_recorder(date,t,share,price,alive,brownian,sigma[0],uniform)
        if np.any(alive&~previous_alive):raise AssertionError('Permanent loss must be absorbing')
        previous_alive=alive;max_loss_error=max(max_loss_error,abs((1-alive.mean())-q))
        max_partition_error=max(max_partition_error,float(np.max(np.abs(share+(1-share)-1))))
        if not j:max_anchor_error=float(np.max(np.abs(price/settings['spot_usd']-1)))
        for threshold,_ in THRESHOLDS:
            ever[threshold]|=(share>=threshold)&alive
            ever_counterfactual[threshold]|=share>=threshold
        selected=(date in target_set or j in (0,len(dates)-1))
        if not targets_only or selected:
            rows,mass=summarize(price,date,cpi,q,uniform,base)
            if not targets_only:monthly+=rows;coverage+=mass
            if date in target_set:targets+=rows
            if selected:
                diagnostics.append(dict(date=date.isoformat(),paths=n,mean=rows[0]['mean'],mean_standard_error=rows[0]['mean_standard_error'],
                    market_log_odds_variance=float((sigma[0]*brownian).var(ddof=1)),expected_market_log_odds_variance=float(sigma[0]**2*diffusion_clock(t,inputs)),
                    realized_log_odds_variance=float(x[0].var(ddof=1)),maximum_root_residual=max_root_residual,
                    loss_probability=q,simulated_loss_probability=float(1-alive.mean())))
        prob=dict(date=date.isoformat(),loss_probability=q,simulated_loss_probability=float(1-alive.mean()),
            nominal_below_initial=float((np.where(alive,price,0)<settings['spot_usd']).mean()),nominal_above_one_million=float(((price>=1e6)&alive).mean()),
            mean_world_wealth_share=float(np.where(alive,share,0).mean()),counterfactual_mean_world_wealth_share=float(share.mean()),
            maximum_simulated_share=float(np.max(np.where(alive,share,0.))),maximum_counterfactual_share=float(share.max()))
        for threshold,label in THRESHOLDS:
            prob['share_above_'+label+'pct']=float(((share>=threshold)&alive).mean())
            prob['counterfactual_share_above_'+label+'pct']=float((share>=threshold).mean())
            prob['ever_share_above_'+label+'pct']=float(ever[threshold].mean())
            prob['counterfactual_ever_share_above_'+label+'pct']=float(ever_counterfactual[threshold].mean())
        for threshold,label in [(.05,'5'),(.10,'10')]:
            prob['mean_contribution_above_'+label+'pct']=float(np.where(alive&(share>=threshold),price,0.).mean()/np.where(alive,price,0.).mean())
        probabilities.append(prob)
        if not targets_only:
            names=[('statistical_only','statistical'),('demand_only','demand'),('equilibrium_only','equilibrium'),('reference','combined')]
            if all(key in index for key,_ in names):
                medians={name:float(np.median(shares[index[key]]))*wealth/quantity for key,name in names}
                for basis,dfl in zip(BASES,[1.,cpi]):components.append(dict(date=date.isoformat(),year=date.year,basis=basis,scope='counterfactual_before_permanent_loss',**{k:v/dfl for k,v in medians.items()}))
            central.append(dict(date=date.isoformat(),wealth_nominal_usd=wealth,cpi_ratio_2026sep=cpi,supply_btc=quantity,
                participating_investable_wealth_share=economic_state(t,settings)['available']))
            moment_rows.append(dict(date=date.isoformat(),market_log_odds_variance=float((sigma[0]*brownian).var(ddof=1)),
                expected_market_log_odds_variance=float(sigma[0]**2*diffusion_clock(t,inputs)),realized_log_odds_variance=float(x[0].var(ddof=1)),
                reference_diffusion=float(sigma[0]*diffusion_scale(t,inputs)),initial_diffusion=float(sigma[0]),
                diffusion_clock_years=float(diffusion_clock(t,inputs)),maximum_root_residual=max_root_residual))
            sample_paths.append(share[:64].copy())
            for key in ['no_loss','loss_low','reference','loss_high','fixed_hazard']:
                c=next((c for c in all_cases if c['case']==key),None)
                if c is None:continue
                risk=float(-np.expm1(-c['hazard']*t)) if c['loss_model']=='constant' else loss.probability(date,scenario=c['loss_scenario'])
                value=float(np.where(uniform>=risk,price,0.).mean())
                for basis,dfl in zip(BASES,[1.,cpi]):loss_curves.append(dict(date=date.isoformat(),case=key,basis=basis,mean=value/dfl))
        if date in target_set:
            for c in all_cases:
                risk=float(-np.expm1(-c['hazard']*t)) if c['loss_model']=='constant' else loss.probability(date,c['loss_model'],c['loss_scenario'])
                k=index[c['process']];candidate=wealth/quantity*shares[k]
                case_rows,_=summarize(candidate,date,cpi,risk,uniform,c)
                for row in case_rows:
                    row['mean_world_wealth_share']=row['mean']*(cpi if row['basis']==BASES[1] else 1)*quantity/wealth
                    row['probability_above_10pct']=float(((shares[k]>=.10)&(uniform>=risk)).mean())
                    row['probability_above_20pct']=float(((shares[k]>=.20)&(uniform>=risk)).mean())
                sensitivity+=case_rows
            all_mean=float(np.where(alive,price,0.).mean());extra_real=float((~alive).mean())*.01*settings['spot_usd']
            for basis,dfl in zip(BASES,[1.,cpi]):
                extra=extra_real*(cpi if basis==BASES[0] else 1.)
                residual.append(dict(year=date.year,basis=basis,loss_probability=q,mean_zero_approximation=all_mean/dfl,
                    mean_with_one_percent_real_residual=all_mean/dfl+extra,relative_mean_change=extra/(all_mean/dfl)))
            print(date.year,'mean',round(all_mean),'share',round(prob['mean_world_wealth_share']*100,3),
                'P20',prob['share_above_20pct'],'ever20',prob['ever_share_above_20pct'],flush=True)
        if j and j%120==0:print(f'  simulated through {date.date()}',flush=True)
    directory.mkdir(parents=True,exist_ok=True)
    for name,rows in [('path_forecast_quantiles.csv',targets),('path_sensitivity_quantiles.csv',sensitivity),('path_sampling_diagnostics.csv',diagnostics),
                      ('permanent_loss_residual_sensitivity.csv',residual),('portfolio_monthly_probabilities.csv',probabilities)]:
        write_csv(directory,name,rows)
    if not targets_only:
        for name,rows in [('path_monthly_quantiles.csv',monthly),('path_interval_probability_mass.csv',coverage),
            ('portfolio_model_monthly_medians.csv',components),('portfolio_macro_central.csv',central),
            ('portfolio_loss_monthly_means.csv',loss_curves),('monetary_monthly_moments.csv',moment_rows)]:write_csv(directory,name,rows)
        np.savez_compressed(directory/'feedback_sample_paths.npz',shares=np.array(sample_paths),times=times)
    (directory/'path_sensitivity_cases.json').write_text(json.dumps(all_cases,ensure_ascii=False,indent=2)+'\n')
    assert max_anchor_error<1e-12 and max_partition_error<1e-12 and max_loss_error<=1/n+1e-12
    result=dict(method=settings['method_version'],paths=n,seed=seed,months=len(dates),simulation_end=dates[-1].isoformat(),
        integration_substeps_per_month=substeps,market_log_odds_diffusion=float(sigma[0]),initial_world_wealth_share=float(s0),
        longrun_market_log_odds_diffusion=float(sigma[0]*diffusion_scale(np.inf,inputs)),
        implied_initial_annual_price_volatility=float(sigma[0]*(1-s0)),equilibrium_prior_retained_fraction=equilibrium.accepted_fraction,
        equilibrium_minimum_terminal_margin=equilibrium.minimum_terminal_margin,
        demand_initial_calibration_error=demand.source_initial_calibration_error,max_anchor_relative_error=max_anchor_error,
        max_partition_sum_error=max_partition_error,max_lifetime_probability_error=max_loss_error,maximum_root_residual=max_root_residual,
        all_checks_passed=True,meaning='Accounting, nonlinear solver and sampling checks; not an empirical validation of feedback coefficients or forecast probabilities',
        hashes={p:hashlib.sha256((ROOT/p).read_bytes()).hexdigest() for p in ['path_assumptions.json','portfolio_inputs.json','portfolio_inputs.npz',
            'permanent_loss_inputs.json','monetary_inputs.json','monetary_inputs.npz','statistical_trend.py','share_feedback.py','portfolio_models.py','macro_path.py','asset_model.py','asset_risk_inputs.json','asset_return_inputs.json']})
    (directory/'path_validation.json').write_text(json.dumps(result,ensure_ascii=False,indent=2)+'\n')
    return result


def main():
    ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('--paths',type=int);ap.add_argument('--seed',type=int)
    ap.add_argument('--output-dir',type=Path,default=ROOT);ap.add_argument('--cache-dir',type=Path)
    ap.add_argument('--targets-only',action='store_true');ap.add_argument('--reference-only',action='store_true')
    ap.add_argument('--substeps',type=int)
    a=ap.parse_args();c=config();result=run(a.paths or c['simulation_paths'],a.seed or c['seed'],a.output_dir,a.cache_dir,a.targets_only,a.reference_only,a.substeps)
    print(json.dumps({k:result[k] for k in ['paths','months','simulation_end','maximum_root_residual','max_lifetime_probability_error']},indent=2))


if __name__=='__main__':main()