from __future__ import annotations
import sys
from itertools import product
from pathlib import Path
import numpy as np
import pandas as pd
sys.path.insert(0,'/mnt/data')
import kartsell_v12_research as base
from kartsell_v12_2_extension import load_kr, KR_FILES, KR_CLASS

OUT=Path('/mnt/data/kartsell_v12_2_results'); OUT.mkdir(exist_ok=True)
US={s:base.load_symbol(s,p) for s,p in base.FILES.items() if s in ['SPY','AAPL','AMZN','NVDA','TSLA','AIG']}
KR={s:load_kr(s,p) for s,p in KR_FILES.items()}
CLASSES={**{s:base.ASSET_CLASS[s] for s in US},**KR_CLASS}
FEATURES={s:base.add_features(d).reset_index(drop=True) for s,d in {**US,**KR}.items()}
BASE_MIN={'broad_index_etf':.35,'large_quality_stock':.25,'high_vol_stock':.20,'cyclical_stock':.20,'distressed_stock':.10}
CRISIS={'broad_index_etf':-.18,'large_quality_stock':-.18,'high_vol_stock':-.22,'cyclical_stock':-.20,'distressed_stock':-.30}

def positions_fast(x, asset_class, act_mult=1., gb_mult=1., wait=10, breakout=20, min_shift=0.):
    close=x.close.to_numpy(float); atr=x.atr20.to_numpy(float); sma50=x.sma50.to_numpy(float); sma200=x.sma200.to_numpy(float)
    slope=x.sma50_slope20.to_numpy(float); vol=x.vol63.to_numpy(float); dates=x.date.to_numpy()
    highprev=pd.Series(close).rolling(breakout,min_periods=max(5,breakout//2)).max().shift(1).to_numpy(float)
    n=len(x); out=np.empty(n,float)
    pos=1.; entry=close[0]; peak=entry; prev_floor=-np.inf; breach=0; last_sell=-10000; last_buy=-10000; below=0
    minpos=float(np.clip(BASE_MIN[asset_class]+min_shift,.05,.75)); act=base.activation_gain(asset_class)*act_mult; crisis=CRISIS[asset_class]
    for i in range(n):
        px=close[i]; peak=max(peak,px); pg=peak/entry-1.; dd=px/peak-1.
        below=below+1 if np.isfinite(sma200[i]) and px<sma200[i] else 0
        regime=np.clip((-dd-.08)/.22,0,1); flow=np.clip((vol[i]-.20)/.50,0,1) if np.isfinite(vol[i]) else 0.
        active=pg>=act or (np.isfinite(atr[i]) and peak-entry>=4*atr[i])
        if active and np.isfinite(atr[i]):
            lock=np.clip(.30+.25*regime+.15*flow,.30,.85); mult=np.clip(3.5-1.7*regime-.8*flow,1.6,3.5)
            gb=np.clip(base.giveback_cap_v11(asset_class,pg)*gb_mult,.05,.40)
            floor=max(entry+lock*(peak-entry),peak-mult*atr[i],peak*(1-gb),prev_floor); prev_floor=floor
            breach=breach+1 if px<floor else 0; gap=(floor-px)/atr[i] if atr[i]>0 else 0
            if i-last_sell>=5 and pos>minpos:
                if gap>=1.5: pos=max(minpos,pos-.40); last_sell=i; breach=0
                elif breach>=2: pos=max(minpos,pos-.20); last_sell=i; breach=0
        if below>=1 and dd<=crisis and i-last_sell>=5 and pos>minpos:
            pos=max(minpos,pos-.25); last_sell=i
        if pos<1 and i-last_sell>=wait and i-last_buy>=wait:
            reclaim=np.isfinite(sma50[i]) and px>sma50[i] and np.isfinite(slope[i]) and slope[i]>0
            br=np.isfinite(highprev[i]) and px>highprev[i]
            if reclaim and br:
                pos=min(1,pos+.20); last_buy=i; entry=px; peak=px; prev_floor=-np.inf; breach=0
        out[i]=pos
    return pd.DataFrame({'date':dates,'position':out})

def bt(s,par):
    x=FEATURES[s]; p=positions_fast(x,CLASSES[s],**par); cost=.0010 if s in KR else .0007
    return base.backtest(x[['date','close']],p,cost=cost)

def market_port(bts,syms,start,end):
    selected={s:base.period_slice(bts[s],start,end) for s in syms}
    selected={s:b for s,b in selected.items() if len(b)>=60}
    return base.portfolio_backtest(selected,False)

def score(m):
    if not m or not np.isfinite(m.get('calmar',np.nan)): return -999.
    return float(m['calmar']-.015*m['turnover_annual']+.25*max(-.3,min(.3,m['cagr'])))

rows=[]
# Compact local-neighbourhood test: activation and giveback ±20%, re-entry frozen.
for act,gb in product([.8,1.,1.2],[.8,1.,1.2]):
    par={'act_mult':act,'gb_mult':gb,'wait':10,'breakout':20,'min_shift':0.}
    bts={s:bt(s,par) for s in FEATURES}; mets={}
    for region,syms,end in [('US',list(US),'2026-03-20'),('KR',list(KR),'2021-04-16')]:
        for split,(st,en) in [('TRAIN',('2007-01-01','2017-12-31')),('OOS',('2018-01-01',end))]:
            p=market_port(bts,syms,st,en); mets[(region,split)]=base.metrics(p) if len(p) else {}
    row={**par,'train_robust_score':min(score(mets[('US','TRAIN')]),score(mets[('KR','TRAIN')])),
         'oos_robust_score':min(score(mets[('US','OOS')]),score(mets[('KR','OOS')]))}
    for (r,sp),m in mets.items():
        for k in ['cagr','mdd','calmar','turnover_annual','avg_exposure']:
            row[f'{r.lower()}_{sp.lower()}_{k}']=m.get(k,np.nan)
    rows.append(row)
res=pd.DataFrame(rows).sort_values('train_robust_score',ascending=False).reset_index(drop=True)
res['train_rank']=np.arange(1,len(res)+1); res['oos_rank']=res.oos_robust_score.rank(ascending=False,method='min')
res.to_csv(OUT/'parameter_robustness_grid_9.csv',index=False)
top=res.head(3)
summary={'variants':len(res),'train_oos_spearman':res[['train_robust_score','oos_robust_score']].corr(method='spearman').iloc[0,1],
         'top3_train_below_median_oos_rate':float((top.oos_robust_score<res.oos_robust_score.median()).mean()),
         'baseline_train_rank':int(res.index[(res.act_mult==1)&(res.gb_mult==1)][0]+1),
         'baseline_oos_rank':int(res.loc[(res.act_mult==1)&(res.gb_mult==1),'oos_rank'].iloc[0])}
pd.DataFrame([summary]).to_csv(OUT/'parameter_robustness_summary_9.csv',index=False)
print(pd.DataFrame([summary]).to_string(index=False)); print(res.to_string(index=False))