"""AHS rates338: educational single-curve Hull–White, no market calibration claim."""
import copy, csv, hashlib, io, json, math, os, re
from datetime import date
import numpy as np
VERSION='338.1'
class InputError(ValueError):
    def __init__(self, field, message, code='INVALID_INPUT'):
        self.field,self.code=field,code
        super().__init__(message)
def number(v, field, lo=-1e100, hi=1e100):
    if isinstance(v,bool) or not isinstance(v,(int,float)) or not math.isfinite(v) or not lo<=v<=hi:
        raise InputError(field,f'{field}: finite number in [{lo}, {hi}] required')
    return float(v)
def integer(v,field,lo,hi):
    number(v,field,lo,hi)
    if int(v)!=v: raise InputError(field,'Integer required')
    return int(v)
def required(p,k):
    if k not in p: raise InputError(k,'Required field missing')
    return p[k]
class Curve:
    def __init__(self,p):
        if set(p)-{'times','discount_factors','interpolation'}: raise InputError('curve','Unknown curve field')
        self.method=required(p,'interpolation')
        if self.method not in ['log_df','linear_df','linear_zero']: raise InputError('interpolation','Unsupported interpolation')
        for key in ['times','discount_factors']:
            arr=required(p,key)
            if not isinstance(arr,list): raise InputError(key,'Array required')
            for v in arr: number(v,key)
        self.times=np.array(required(p,'times'),dtype=float)
        self.dfs=np.array(required(p,'discount_factors'),dtype=float)
        if len(self.times)<2 or len(self.times)!=len(self.dfs) or not np.isfinite(self.times).all() or not np.isfinite(self.dfs).all() or self.times[0]!=0 or abs(self.dfs[0]-1)>1e-12 or (np.diff(self.times)<=0).any() or (self.dfs<=0).any():
            raise InputError('curve','Require increasing times starting at 0, matching positive DFs and P(0)=1')
        self.zeros=np.r_[-math.log(self.dfs[1])/self.times[1],-np.log(self.dfs[1:])/self.times[1:]]
    def df(self,t):
        if t<0 or t>self.times[-1]+1e-10: raise InputError('time','Extrapolation disabled')
        if self.method=='log_df': return math.exp(float(np.interp(t,self.times,np.log(self.dfs))))
        if self.method=='linear_df': return float(np.interp(t,self.times,self.dfs))
        return math.exp(-t*float(np.interp(t,self.times,self.zeros)))
    def forward(self,t):
        h=1e-5; l=max(0,t-h); r=min(self.times[-1],t+h)
        return -(math.log(self.df(r))-math.log(self.df(l)))/(r-l)
def ou_moments(a,sigma,dt):
    if os.environ.get('AHS_RATES338_NATIVE')=='1':
        from .native import moments
        return moments(a,sigma,dt)
    decay=math.exp(-a*dt)
    variance=sigma*sigma*(-math.expm1(-2*a*dt))/(2*a)
    return decay,variance
class Tree:
    """Moment-matching recombining OU tree; state-price fitted per time step."""
    def __init__(self,curve,a,sigma,maturity,steps):
        self.dt=maturity/steps; self.steps=steps
        decay,var=ou_moments(a,sigma,self.dt)
        self.dx=math.sqrt(3*var) if var>0 else 0.
        self.nodes=[{0:1.}]; self.transitions=[]; self.discounts=[]; self.fit=[]
        for i in range(steps):
            q=self.nodes[-1]; trans={}; disc={}
            raw=sum(v*math.exp(-j*self.dx*self.dt) for j,v in q.items())
            phi=math.log(raw/curve.df((i+1)*self.dt))/self.dt
            nxt={}
            for j,v in q.items():
                mu=j*decay; k=math.floor(mu+0.5); e=mu-k
                tr=[(k-1,(1/3+e*e-e)/2),(k,2/3-e*e),(k+1,(1/3+e*e+e)/2)] if self.dx else [(0,1.)]
                trans[j]=tr; disc[j]=math.exp(-(j*self.dx+phi)*self.dt)
                for dest,p in tr: nxt[dest]=nxt.get(dest,0.)+v*disc[j]*p
            self.transitions.append(trans); self.discounts.append(disc); self.nodes.append(nxt)
            self.fit.append(sum(nxt.values())-curve.df((i+1)*self.dt))
    def back(self,i,v):
        return {j:self.discounts[i][j]*sum(p*v[k] for k,p in tr) for j,tr in self.transitions[i].items()}
    def bond(self,i,end):
        v={j:1. for j in self.nodes[end]}
        for k in range(end-1,i-1,-1): v=self.back(k,v)
        return v
    def price(self,trade,exercises):
        def idx(t):
            f=t/self.dt; n=round(f)
            if abs(n-f)>1e-7: raise InputError('steps','All payments and exercise times must align with grid')
            return n
        payment=[idx(t) for t in trade['payment_times']]; exercise={idx(t) for t in exercises}
        v={j:0. for j in self.nodes[-1]}; policy=[]
        for i in range(self.steps-1,-1,-1):
            v=self.back(i,v)
            if i in exercise:
                bonds=[(t,self.bond(i,k)) for t,k in zip(trade['payment_times'],payment) if k>i]
                end=bonds[-1][1]; prev=i*self.dt
                fixed={j:0. for j in self.nodes[i]}
                for t,b in bonds:
                    for j in fixed: fixed[j]+=trade['fixed_rate']*(t-prev)*b[j]
                    prev=t
                for j in v:
                    immediate=max(trade['direction']*trade['notional']*(1-end[j]-fixed[j]),0.)
                    policy.append([i*self.dt,j*self.dx,immediate,v[j],int(immediate>v[j])])
                    v[j]=max(v[j],immediate)
        return v[0],policy

def swap_pv(curve,tr):
    fixed=0.; prev=tr['start']
    for t in tr['payment_times']:
        fixed+=(t-prev)*curve.df(t)*tr['fixed_rate']*tr['notional']; prev=t
    floating=tr['notional']*(curve.df(tr['start'])-curve.df(tr['payment_times'][-1]))
    return {'fixed_leg':fixed,'floating_leg':floating,'price':tr['direction']*(floating-fixed)}

def mc_bond(curve,a,sigma,T,paths,seed,steps):
    """Exact joint OU state/integral; deterministic shift integrated from initial curve."""
    rng=np.random.default_rng(seed); x=np.zeros(paths); integral=np.zeros(paths); rows=[]; dt=T/steps
    d,v=ou_moments(a,sigma,dt); B=(1-d)/a
    cov=sigma*sigma*(1-d)**2/(2*a*a)
    vi=sigma*sigma/a**2*(dt-2*(1-d)/a+(1-d*d)/(2*a))
    vi=max(vi,0.)
    for k in range(steps):
        z=rng.standard_normal((2,paths)); old=x.copy()
        xi=math.sqrt(v)*z[0] if v else np.zeros(paths)
        eta=(cov/math.sqrt(v)*z[0]+math.sqrt(max(vi-cov*cov/v,0))*z[1]) if v else np.zeros(paths)
        integral+=B*old+eta; x=d*old+xi
        t=(k+1)*dt; shift=curve.forward(t)+sigma*sigma/(2*a*a)*(1-math.exp(-a*t))**2
        rows.append([t,float(x[0]+shift),float(np.mean(x)),float(np.var(x)),sigma*sigma*(1-math.exp(-2*a*t))/(2*a)])
    varI=sigma*sigma/a**2*(T-2*(1-math.exp(-a*T))/a+(1-math.exp(-2*a*T))/(2*a))
    samples=curve.df(T)*np.exp(-integral-.5*varI)
    convergence=[]
    for n in sorted(set([min(paths,n) for n in [64,256,1024,4096,paths]])):
        mean=float(np.mean(samples[:n])); se=float(np.std(samples[:n],ddof=1)/math.sqrt(n)); convergence.append([n,mean,se,mean-1.96*se,mean+1.96*se,curve.df(T)])
    return {'convergence':convergence,'moments':rows,'price':float(np.mean(samples)),'standard_error':float(np.std(samples,ddof=1)/math.sqrt(paths)),'reference':curve.df(T),'discretisation':'Exact Gaussian OU integral; no Euler bias for this zero-coupon payoff','seed':seed}

def tarn_path(rates,p):
    accrued=0.; rows=[]; active=True
    for i,r in enumerate(rates):
        raw=max(p['floor'],p['coupon']+p['gearing']*r)*p['accrual']
        paid=min(raw,max(p['target']-accrued,0)) if active else 0.
        accrued+=paid
        if accrued>=p['target']-1e-12: active=False
        rows.append([i+1,r,raw,paid,accrued,int(active)])
    return rows

class Registry:
    def __init__(self): self.objects={}
    def put(self,kind,data):
        payload={'kind':kind,'data':copy.deepcopy(data),'engine':VERSION}
        key=kind+'_'+hashlib.sha256(json.dumps(payload,sort_keys=True,allow_nan=False).encode()).hexdigest()[:20]
        self.objects[key]=payload; return {'id':key,'kind':kind,'engine':VERSION}
    def get(self,key,kind):
        if key not in self.objects: raise InputError('id','Unknown ID; process may have restarted','NOT_FOUND')
        if self.objects[key]['kind']!=kind: raise InputError('id','Object type mismatch','TYPE_MISMATCH')
        return copy.deepcopy(self.objects[key]['data'])
    def call(self,op,p):
        try: return self._call(op,p)
        except InputError: raise
        except (ValueError,TypeError,KeyError,OverflowError,IndexError) as e:
            raise InputError('payload','Malformed input type, dimensions, missing field or numeric range') from e
    def _call(self,op,p):
        if not isinstance(p,dict): raise InputError('payload','Object required')
        fields={
          'CreateMarket':{'valuation_date','currency'},'AddCurve':{'market_id','curve'},'AddVolatilitySurface':{'market_id','surface'},'SetModelParameters':{'market_id','a','sigma'},
          'CreateSwap':{'notional','start','payment_times','fixed_rate','direction'},'CreateExerciseSchedule':{'times'},'CreateBermudanSwaption':{'swap_id','exercise_id'},
          'CreateTARN':{'notional','rates','target','coupon','gearing','floor','accrual'},'CreatePricer':{'market_id','trade_id','steps'},'RunValuation':{'pricer_id'},
          'ComputeSensitivities':{'pricer_id','bump'},'RunScenario':{'pricer_id','parallel_shift'},'GetDiagnostics':{'pricer_id'},'ExportResults':{'result_id'}}
        if op in fields and set(p)-fields[op]: raise InputError(sorted(set(p)-fields[op])[0],'Unknown field')

        if op=='CreateMarket':
            d=required(p,'valuation_date')
            if not isinstance(d,str) or not re.fullmatch(r'\d{4}-\d{2}-\d{2}',d): raise InputError('valuation_date','ISO date YYYY-MM-DD required')
            try: date.fromisoformat(d)
            except (ValueError,TypeError): raise InputError('valuation_date','ISO date YYYY-MM-DD required')
            if required(p,'currency')!='EUR': raise InputError('currency','Only EUR educational market supported')
            return self.put('market',{'valuation_date':d,'currency':'EUR'})
        if op in ['AddCurve','AddVolatilitySurface','SetModelParameters']:
            m=self.get(required(p,'market_id'),'market')
            if op=='AddCurve':
                c=required(p,'curve'); Curve(c); m['curve']=c
            if op=='AddVolatilitySurface':
                s=required(p,'surface')
                if set(s)-{'expiries','tenors','values','quote_convention'}: raise InputError('surface','Unknown surface field')
                for row in s['values']:
                    for v in row: number(v,'values',0,100)
                if required(s,'quote_convention')!='normal_decimal_sqrt_year': raise InputError('quote_convention','Only stored normal decimal/sqrt(year) quotes supported')
                for k in ['expiries','tenors']:
                    a=s[k]
                    if not a or any(number(x,k,1e-8,100)<=0 for x in a) or any(a[i]>=a[i+1] for i in range(len(a)-1)): raise InputError(k,'Strictly increasing positive axis required')
                vals=np.array(s['values'],float)
                if vals.shape!=(len(s['expiries']),len(s['tenors'])) or not np.isfinite(vals).all() or (vals<0).any(): raise InputError('values','Invalid volatility matrix')
                m['surface']=s
            if op=='SetModelParameters':
                m['model']={'a':number(required(p,'a'),'a',.001,3),'sigma':number(required(p,'sigma'),'sigma',0,.1),'calibration':'User supplied; surface stored only, not calibrated'}
            return self.put('market',m)
        if op=='CreateSwap':
            notional=number(required(p,'notional'),'notional',1,1e10); start=number(required(p,'start'),'start',0,30)
            times=required(p,'payment_times')
            if not isinstance(times,list) or not times or len(times)>60: raise InputError('payment_times','1–60 payments required')
            prev=start
            for t in times:
                number(t,'payment_times',prev+1e-8,30); prev=t
            direction=required(p,'direction')
            if direction not in [-1,1] or isinstance(direction,bool): raise InputError('direction','1 payer fixed; -1 receiver fixed')
            return self.put('swap',dict(notional=notional,start=start,payment_times=times,fixed_rate=number(required(p,'fixed_rate'),'fixed_rate',-.2,.5),direction=direction))
        if op=='CreateExerciseSchedule':
            times=required(p,'times')
            if not isinstance(times,list) or not times: raise InputError('times','Nonempty list required')
            prev=0.
            for t in times: number(t,'times',prev+1e-8,30); prev=t
            return self.put('exercise',{'times':times})
        if op=='CreateBermudanSwaption':
            tr=self.get(required(p,'swap_id'),'swap'); ex=self.get(required(p,'exercise_id'),'exercise')['times']
            if any(t<tr['start'] or t>=tr['payment_times'][-1] or (t!=tr['start'] and t not in tr['payment_times']) for t in ex): raise InputError('exercise','Exercise must align with start/payment boundaries before maturity')
            return self.put('bermudan',{'swap_id':p['swap_id'],'exercise_id':p['exercise_id'],'settlement':'Physical: enter remaining co-terminal swap at exercise; zero fee; zero notice'})
        if op=='CreateTARN':
            data={k:number(required(p,k),k,*bounds) for k,bounds in {'target':(1e-8,1),'coupon':(-1,1),'gearing':(-20,20),'floor':(0,1),'accrual':(.01,1),'notional':(1,1e10)}.items()}
            rates=required(p,'rates')
            if not isinstance(rates,list) or not rates or len(rates)>120: raise InputError('rates','1–120 rates required')
            data['rates']=[number(r,'rates',-.5,1) for r in rates]
            return self.put('tarn',data)
        if op=='CreatePricer':
            m=self.get(required(p,'market_id'),'market')
            if 'curve' not in m or 'model' not in m: raise InputError('market_id','Curve and explicit model parameters required')
            key=required(p,'trade_id')
            if key not in self.objects or self.objects[key]['kind'] not in ['swap','bermudan','tarn']: raise InputError('trade_id','Supported trade ID required')
            steps=integer(required(p,'steps'),'steps',4,600)
            config={'market_id':p['market_id'],'trade_id':key,'steps':steps}
            # Validate coverage and grid before assigning an executable pricer ID.
            self.value(config)
            return self.put('pricer',config)
        if op=='RunValuation':
            config=self.get(required(p,'pricer_id'),'pricer'); result=self.value(config)
            result.update({'pricer_id':p['pricer_id'],'configuration':config,'market_snapshot':self.get(config['market_id'],'market'),'trade_snapshot':copy.deepcopy(self.objects[config['trade_id']]),'engine':VERSION,'backend':'Python + C++ OU kernel' if os.environ.get('AHS_RATES338_NATIVE')=='1' else 'Python/NumPy'})
            result_id=self.put('result',result)['id']; return {'result_id':result_id,**result}
        if op=='ComputeSensitivities':
            cfg=self.get(required(p,'pricer_id'),'pricer'); bump=number(required(p,'bump'),'bump',1e-6,.01)
            mid=self.value(cfg)['price']; up=self.value(cfg,rate_shift=bump)['price']; down=self.value(cfg,rate_shift=-bump)['price']
            return {'base':mid,'up':up,'down':down,'central_change':(up-down)/2,'derivative_per_decimal_rate':(up-down)/(2*bump),'bump_decimal':bump,'definition':'Parallel continuous zero-rate shift; model parameters fixed, tree refitted; + = value rises when rates rise','unit':'EUR'}
        if op=='RunScenario':
            cfg=self.get(required(p,'pricer_id'),'pricer'); shift=number(required(p,'parallel_shift'),'parallel_shift',-.1,.1)
            return {'parallel_shift':shift,**self.value(cfg,shift)}
        if op=='GetDiagnostics':
            cfg=self.get(required(p,'pricer_id'),'pricer'); base=self.value(cfg)
            levels=sorted(set([max(4,cfg['steps']//2),cfg['steps'],min(600,cfg['steps']*2)])); out=[]
            for n in levels:
                try: out.append({'steps':n,'price':self.value({**cfg,'steps':n})['price']})
                except InputError as e: out.append({'steps':n,'error':str(e)})
            return {'curve_fit_max_abs':base.get('curve_fit_max_abs'),'grid_comparison':out,'validation':'Runtime diagnostics only; grid differences are not a proof of convergence','scope':'Single curve EUR; year fractions; no calendars, fixings, collateral, XVA or automatic vol calibration'}
        if op=='ExportResults':
            r=self.get(required(p,'result_id'),'result'); f=io.StringIO(); w=csv.writer(f); w.writerow(['field','value','unit'])
            for k,v in r.items(): w.writerow([k,json.dumps(v,ensure_ascii=False,allow_nan=False) if isinstance(v,(dict,list)) else v,'EUR' if k=='price' else ''])
            return {'filename':'ahs-rates338.csv','csv':f.getvalue()}
        raise InputError('operation','Unknown operation','UNKNOWN_OPERATION')
    def value(self,cfg,rate_shift=0):
        m=self.get(cfg['market_id'],'market'); c=copy.deepcopy(m['curve']); c['discount_factors']=[d*math.exp(-rate_shift*t) for d,t in zip(c['discount_factors'],c['times'])]; curve=Curve(c)
        item=self.objects[cfg['trade_id']]; kind=item['kind']; tr=copy.deepcopy(item['data'])
        if kind=='swap': return {**swap_pv(curve,tr),'currency':'EUR','method':'Discounted cash flows, single curve'}
        if kind=='tarn':
            rows=tarn_path(tr['rates'],tr); pv=sum(row[3]*curve.df(row[0]*tr['accrual'])*tr['notional'] for row in rows)
            return {'price':pv,'currency':'EUR','path':rows,'method':'Exotic coupon leg on user-supplied deterministic rate scenario; not a stochastic full-TARN price','rules':'Coupon capped at remaining target; no maturity top-up; funding leg excluded; no principal exchange'}
        swap=self.get(tr['swap_id'],'swap'); ex=self.get(tr['exercise_id'],'exercise')['times']; a=m['model']['a']; sigma=m['model']['sigma']
        tree=Tree(curve,a,sigma,swap['payment_times'][-1],cfg['steps']); pv,policy=tree.price(swap,ex)
        return {'price':pv,'currency':'EUR','method':'Hull–White fitted trinomial tree','curve_fit_max_abs':max(abs(x) for x in tree.fit),'exercise_policy':policy,'grid':{'steps':cfg['steps'],'dt_years':tree.dt,'dx_rate_decimal':tree.dx,'max_nodes':max(map(len,tree.nodes)),'boundary':'Natural recombining support; no imposed clipping; inspect grid convergence'}}

def integral_variance(a,sigma,t):
    return sigma*sigma/a**2*(t-2*(-math.expm1(-a*t))/a+(-math.expm1(-2*a*t))/(2*a))
def conditional_bond(curve,a,sigma,t,T,x):
    b=-math.expm1(-a*(T-t))/a
    return curve.df(T)/curve.df(t)*np.exp(-b*x+.5*(integral_variance(a,sigma,T-t)-integral_variance(a,sigma,T)+integral_variance(a,sigma,t)))
def regression_bermudan(curve,trade,exercise,a,sigma,paths=4096,seed=338,degree=2):
    """Continuation LSM, separate train/evaluation draws. Not pickup regression."""
    paths=integer(paths,'paths',128,20000); degree=integer(degree,'degree',1,3)
    rng=np.random.default_rng(seed)
    def simulate():
        x=np.zeros(paths); integ=np.zeros(paths); xs=[]; ds=[]; last=0.
        for t in exercise:
            dt=t-last; d,v=ou_moments(a,sigma,dt); b=(1-d)/a
            cov=sigma*sigma*(1-d)**2/(2*a*a); vi=max(integral_variance(a,sigma,dt),0.)
            z=rng.standard_normal((2,paths)); xi=math.sqrt(v)*z[0]
            eta=(cov/math.sqrt(v)*z[0]+math.sqrt(max(vi-cov*cov/v,0))*z[1]) if v else np.zeros(paths)
            integ+=b*x+eta; x=d*x+xi
            xs.append(x.copy()); ds.append(curve.df(t)*np.exp(-integ-.5*integral_variance(a,sigma,t))); last=t
        return np.array(xs),np.array(ds)
    def payoffs(xs):
        out=[]
        for i,t in enumerate(exercise):
            fixed=np.zeros(paths); prev=t
            for T in trade['payment_times']:
                if T>t:
                    fixed+=trade['fixed_rate']*(T-prev)*conditional_bond(curve,a,sigma,t,T,xs[i]); prev=T
            end=conditional_bond(curve,a,sigma,t,trade['payment_times'][-1],xs[i])
            out.append(np.maximum(trade['direction']*trade['notional']*(1-end-fixed),0))
        return np.array(out)
    x,d=simulate(); h=payoffs(x); cash=h[-1]*d[-1]; coeff={}; ranks=[]
    def design(v,t):
        scale=max(math.sqrt(ou_moments(a,sigma,t)[1]),1e-6)
        return np.array([(v/scale)**k for k in range(degree+1)]).T
    for i in range(len(exercise)-2,-1,-1):
        X=design(x[i],exercise[i]); mask=h[i]>0
        if np.count_nonzero(mask)>=degree+2:
            co,_,rank,_=np.linalg.lstsq(X[mask],cash[mask]/d[i,mask],rcond=None)
        else: co=np.zeros(degree+1); co[0]=float(np.mean(cash/d[i])); rank=0
        coeff[i]=co; ranks.append({'time':exercise[i],'rank':int(rank),'itm_paths':int(mask.sum()),'basis':'1, x/std(x), ..., power '+str(degree)})
        cont=np.maximum(X@co,0); act=(h[i]>cont)&mask; cash=np.where(act,h[i]*d[i],cash)
    train=float(np.mean(cash)); xe,de=simulate(); he=payoffs(xe); alive=np.ones(paths,bool); value=np.zeros(paths); policy=[]
    for i,t in enumerate(exercise):
        cont=np.maximum(design(xe[i],t)@coeff[i],0) if i in coeff else np.zeros(paths)
        act=alive&(he[i]>0)&((he[i]>=cont) if i in coeff else np.ones(paths,bool))
        for k in range(min(paths,160)): policy.append([t,float(xe[i,k]),float(he[i,k]),float(cont[k]),int(act[k])])
        value+=np.where(act,he[i]*de[i],0); alive&=~act
    return {'price':float(np.mean(value)),'standard_error':float(np.std(value,ddof=1)/math.sqrt(paths)),'training_price':train,'train_paths':paths,'evaluation_paths':paths,'degree':degree,'seed':seed,'policy':policy,'regression':ranks,'method':'LSM continuation regression; independent train/evaluation draws; not appendix-D pickup regression','limitations':'Finite basis/finite sample stopping policy; held-out SE excludes policy approximation error'}
