"""Causal ridge ensemble with shrunk stock-specific risk estimates.

CTF submission module: no file I/O, no network calls, no external data.
Algorithm fixed before inspecting competition test-period performance.
"""
from collections import deque

import numpy as np
import pandas as pd


def feature_names(features):
    if 'features' not in features:
        raise ValueError('features must contain the official features column')
    return sorted(features['features'].dropna().astype(str).unique().tolist())


def test_flag(values):
    text = values.astype(str).str.lower().str.strip()
    if not text.isin(['0', '1', '0.0', '1.0', 'false', 'true', 'f', 't']).all():
        raise ValueError('Unexpected ctff_test values')
    return text.isin(['1', '1.0', 'true', 't']).to_numpy()


def transform(frame, names):
    # Monthly cross-section only; no fitted quantities from future months.
    raw = frame[names].astype(np.float32).replace([np.inf, -np.inf], np.nan)
    x = (2 * raw.rank(method='average', pct=True) - 1).fillna(0).to_numpy(dtype=np.float64)
    x -= x.mean(axis=0)
    return x


def rank_score(prediction):
    score = pd.Series(prediction).rank(method='average').to_numpy(dtype=float)
    score -= score.mean()
    scale = np.max(np.abs(score))
    return score / scale if scale > 0 else np.zeros_like(score)


def unit_gross(score):
    gross = np.abs(score).sum()
    return score * (2 / gross) if gross > 1e-12 else np.zeros_like(score)


def fit_ridge(a, b, months):
    if months == 0:
        return np.zeros(len(b))
    # Objective: equally weighted monthly average squared error + L2 penalty.
    penalty = 0.10 * max(np.trace(a) / len(b), 1e-12)
    return np.linalg.solve(a + np.eye(len(b)) * penalty, b)


def risk_weights(ids, score, history, date):
    cutoff = date - pd.DateOffset(months=60)
    beta = np.ones(len(ids))
    variance = np.full(len(ids), np.nan)
    for j, stock in enumerate(ids):
        rows = [v for v in history.get(int(stock), ()) if v[0] > cutoff]
        if len(rows) < 12:
            continue
        observations = np.asarray([[v[1], v[2]] for v in rows], dtype=float)
        r = observations[:, 0]
        m = observations[:, 1]
        r -= r.mean()
        m -= m.mean()
        market_var = m @ m
        raw_beta = (r @ m) / market_var if market_var > 1e-12 else 1.0
        shrink = len(rows) / (len(rows) + 24.0)
        beta[j] = shrink * raw_beta + (1 - shrink)
        variance[j] = np.mean((r - raw_beta * m) ** 2)
    finite = variance[np.isfinite(variance) & (variance > 0)]
    median = float(np.median(finite)) if len(finite) else 0.01
    variance = np.where(np.isfinite(variance), variance, median)
    variance = np.maximum(0.5 * variance + 0.5 * median, median * 0.25)
    precision = 1.0 / variance
    # Shrunk one-factor covariance: market exposure constrained to zero;
    # idiosyncratic variances govern the remaining allocation.
    centered_beta = beta - beta.mean()
    constraints = np.column_stack([np.ones(len(ids)), centered_beta])
    gram = constraints.T @ (precision[:, None] * constraints)
    adjustment = np.linalg.pinv(gram, rcond=1e-10) @ (constraints.T @ (precision * score))
    raw = precision * (score - constraints @ adjustment)
    return unit_gross(raw), beta


def run_months(months, names, diagnostics=False):
    np.random.seed(42)
    dimension = len(names)
    a = np.zeros((dimension, dimension))
    b = np.zeros(dimension)
    recent_a = a.copy()
    recent_b = b.copy()
    realized = deque()
    pending = deque()
    history = {}
    count = 0
    fitted_year = None
    coef_full = np.zeros(dimension)
    coef_recent = np.zeros(dimension)
    outputs = []
    validation_outputs = []
    seen_test_rows = False
    records = []
    importance = []
    for date, frame in months:
        date = pd.Timestamp(date)
        frame = frame.sort_values('id', kind='stable').reset_index(drop=True)
        ids_float = frame['id'].to_numpy(dtype=float)
        if not np.isfinite(ids_float).all() or not np.equal(ids_float, np.floor(ids_float)).all():
            raise ValueError('Security identifiers must be finite integers')
        ids = ids_float.astype(np.int64)
        if len(np.unique(ids)) != len(ids):
            raise ValueError('Duplicate security identifiers within a month')
        # Release labels ONLY on their documented realization date.
        while pending and pending[0][0] <= date:
            available, moment_a, moment_b, past_ids, returns = pending.popleft()
            a += moment_a
            b += moment_b
            recent_a += moment_a
            recent_b += moment_b
            count += 1
            realized.append((available, moment_a, moment_b))
            market = float(np.mean(returns))
            for stock, r in zip(past_ids, returns):
                entries = history.setdefault(int(stock), deque(maxlen=60))
                entries.append((available, float(r), market))
        cutoff = date - pd.DateOffset(months=120)
        while realized and realized[0][0] <= cutoff:
            _, old_a, old_b = realized.popleft()
            recent_a -= old_a
            recent_b -= old_b
        if fitted_year != date.year and count:
            coef_full = fit_ridge(a, b, count)
            coef_recent = fit_ridge(recent_a, recent_b, len(realized))
            fitted_year = date.year
            if diagnostics:
                print(f'{date.date()}: refit using {count} realized months; recent={len(realized)}', flush=True)
                importance.append({'eom': str(date.date()), 'features': names,
                    'coef_expanding': coef_full.tolist(), 'coef_recent': coef_recent.tolist()})
        x = transform(frame, names)
        full_score = rank_score(x @ coef_full)
        recent_score = rank_score(x @ coef_recent)
        ensemble_score = 0.5 * (full_score + recent_score)
        baseline = unit_gross(full_score)
        equal_risk = unit_gross(ensemble_score)
        weights, beta = risk_weights(ids, ensemble_score, history, date)
        selected = test_flag(frame['ctff_test'])
        if selected.any():
            seen_test_rows = True
            validation_outputs.clear()
            outputs.append(pd.DataFrame({'id': ids[selected], 'eom': date, 'w': weights[selected]}))
        elif not seen_test_rows:
            # An all-pre-test validation panel must still execute the real model
            # and return a non-empty schema-valid table (CTF Rules 12 and 14).
            validation_outputs.append(pd.DataFrame({'id': ids, 'eom': date, 'w': weights}))
        # Current y is used only for diagnostic evaluation and queued training;
        # it cannot influence the weights just computed.
        y = frame['ret_exc_lead1m'].to_numpy(dtype=float)
        valid = np.isfinite(y)
        if diagnostics:
            records.append({'eom': str(date.date()), 'eom_ret': str((date + pd.offsets.MonthEnd(1)).date()),
                'test': bool(selected.any()), 'n': len(ids), 'trained_months': count,
                'baseline_return': float(baseline[valid] @ y[valid]),
                'ensemble_return': float(equal_risk[valid] @ y[valid]),
                'model_return': float(weights[valid] @ y[valid]),
                'market_return': float(np.mean(y[valid])) if valid.any() else 0,
                'gross': float(np.abs(weights).sum()), 'net': float(weights.sum()),
                'beta_exposure': float(weights @ beta), 'max_weight': float(np.abs(weights).max()),
                'label_coverage': float(valid.mean())})
        if valid.any():
            dates = pd.to_datetime(frame.loc[valid, 'eom_ret']) if 'eom_ret' in frame else pd.Series(date + pd.offsets.MonthEnd(1), index=np.flatnonzero(valid))
            if (dates <= date).any():
                raise ValueError('Future-return labels must realize after their formation date')
            # Official data realizes every monthly label at the same next month-end.
            if dates.nunique() != 1:
                raise ValueError('Inconsistent realization dates within a month')
            lower, upper = np.quantile(y[valid], [0.01, 0.99])
            target = np.clip(y[valid], lower, upper)
            target -= target.mean()
            train_x = x[valid].copy()
            train_x -= train_x.mean(axis=0)
            moment_a = train_x.T @ train_x / len(target)
            moment_b = train_x.T @ target / len(target)
            pending.append((pd.Timestamp(dates.iloc[0]), moment_a, moment_b, ids[valid], y[valid]))
    frames = outputs if seen_test_rows else validation_outputs
    result = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame(columns=['id', 'eom', 'w'])
    if diagnostics:
        return result, pd.DataFrame(records), importance
    return result


def _original_main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    """Return monthly weights; daily_ret is accepted but not needed by this model."""
    names = feature_names(features)
    required = ['id', 'eom', 'ctff_test', 'ret_exc_lead1m'] + names
    absent = sorted(set(required) - set(chars.columns))
    if absent:
        raise ValueError('Missing required input columns: ' + ', '.join(absent))
    dates = pd.to_datetime(chars['eom'])
    groups = dates.groupby(dates, sort=True).indices
    return run_months(((date, chars.iloc[idx]) for date, idx in groups.items()), names)


# Research attribution: Geertsema's raw-feature LightGBM, bounded target and
# shrunk daily covariance; Keshri's beta-neutral risk construction. Our code
# uses a time-dimension covariance solve and fixed pre-test candidate selection.
# The 1990-2023 historical sample was already viewed for version 1.

"""Original implementation of upgraded components, inspired by cited public research."""
import numpy as np
import pandas as pd
import lightgbm as lgb
from scipy.linalg import cho_factor, cho_solve

PARAMS={'objective':'regression','learning_rate':.02,'num_leaves':63,
 'min_data_in_leaf':2000,'lambda_l2':1.0,'max_bin':63,'num_threads':4,
 'seed':42,'feature_fraction_seed':42,'bagging_seed':42,'data_random_seed':42,
 'deterministic':True,'force_col_wise':True,'verbosity':-1}
ROUNDS=500

def standardized_signal(values):
    values=np.nan_to_num(np.asarray(values,dtype=float),nan=0,posinf=0,neginf=0)
    values-=values.mean()
    scale=values.std()
    return np.clip(values/scale,-3,3) if scale>1e-12 else np.zeros_like(values)

def target_transform(meta):
    result=np.full(len(meta),np.nan,dtype=np.float32)
    for _,idx in meta.groupby('eom',sort=True).indices.items():
        y=meta.iloc[idx].ret_exc_lead1m.to_numpy(dtype=float)
        valid=np.isfinite(y)
        if valid.sum()<2:continue
        scale=np.std(y[valid],ddof=1)
        result[idx[valid]]=np.arctan((y[valid]-y[valid].mean())/max(scale,1e-12))
    return result

def fit_tree(raw,meta,target,date,callbacks=None):
    # Fit at a calendar-defined origin: skipping pre-test predictions must not
    # change which training observations a later forecast uses.
    cutoff=min(pd.Timestamp(date),pd.Timestamp(date.year,1,31))
    use=(pd.to_datetime(meta.eom_ret)<=cutoff).to_numpy() & np.isfinite(target)
    if meta.loc[use,'eom'].nunique()<36:
        return None,cutoff
    params=PARAMS.copy();params['min_data_in_leaf']=min(2000,max(20,int(use.sum())//50))
    # Chronologically sorted, complete panels have a contiguous training prefix.
    # Retain a view of that prefix instead of duplicating the feature matrix.
    count=int(use.sum())
    prefix=bool(np.all(use[:count]) and not np.any(use[count:]))
    x=np.asarray(raw[:count] if prefix else raw[use],dtype=np.float32)
    labels=target[:count] if prefix else target[use]
    dataset=lgb.Dataset(x,label=labels,free_raw_data=True)
    tree=lgb.train(params,dataset,num_boost_round=ROUNDS,callbacks=callbacks)
    return tree,cutoff

def daily_covariance_weights(matrix,dates,daily_ids,ids,signal,date):
    signal=np.asarray(signal,dtype=float)
    single=signal.ndim==1
    if single:signal=signal[:,None]
    ns=signal.shape[1]
    end=int(np.searchsorted(dates,np.datetime64(date,'D'),side='right'))
    start=max(0,end-756)
    cols=np.searchsorted(daily_ids,ids)
    present=cols<len(daily_ids)
    present[present]=daily_ids[cols[present]]==ids[present]
    available=np.flatnonzero(present)
    raw=np.asarray(matrix[start:end,cols[present]],dtype=np.float64)
    enough=np.isfinite(raw).sum(axis=0)>=250
    positions=available[enough]
    if len(positions)<5 or end-start<2:
        return None,{'held':0,'delta':None,'beta_exposure':None}
    x=np.where(np.isfinite(raw[:,enough]),raw[:,enough],0.)
    x-=x.mean(axis=0)
    t,n=x.shape
    # Exact Ledoit-Wolf covariance calculation in the smaller time dimension.
    # ||X'X/T||_F^2 == ||XX'/T||_F^2; no N-by-N covariance is needed.
    gram=(x@x.T)/t
    trace=float(np.trace(gram))
    norm=float(np.square(gram).sum())
    distance=max(norm-trace*trace/n,0.)
    row_norm=np.einsum('ij,ij->i',x,x)
    noise=max(float(row_norm@row_norm)/(t*t)-norm/t,0.)
    delta=min(noise,distance)/distance if distance>1e-20 else 1.
    diagonal=(delta+1.)*trace/n
    if diagonal<=1e-12:
        return None,{'held':0,'delta':delta,'beta_exposure':None}
    coefficient=1.-delta
    market=x.mean(axis=1)
    beta=(x.T@market)/max(float(market@market),1e-12)
    beta_centered=beta-beta.mean()
    constraints=np.column_stack([np.ones(n),beta_centered])
    rhs=np.column_stack([signal[positions],constraints])
    factor=cho_factor(np.eye(t)+(coefficient/diagonal)*gram,lower=True,check_finite=False)
    inv_rhs=rhs/diagonal-(coefficient/(diagonal*diagonal*t))*(x.T@cho_solve(factor,x@rhs,check_finite=False))
    inv_mu=inv_rhs[:,:ns];inv_c=inv_rhs[:,ns:]
    adjustment=np.linalg.pinv(constraints.T@inv_c,rcond=1e-10)@(constraints.T@inv_mu)
    held=inv_mu-inv_c@adjustment
    gross=np.abs(held).sum(axis=0)
    result=np.zeros((len(ids),ns))
    result[positions]=np.divide(held*2,gross,out=np.zeros_like(held),where=gross>1e-12)
    info={'held':n,'delta':float(delta),'beta_exposure':(result[positions].T@beta).tolist()}
    return result[:,0] if single else result,info

SELECTED = 'blend_daily'


def _prepare_monthly(chars,features):
    names=feature_names(features)
    keep=['id','eom','ctff_test','ret_exc_lead1m']
    if 'eom_ret' in chars:keep.insert(2,'eom_ret')
    meta=chars[keep].copy()
    meta['eom']=pd.to_datetime(meta.eom)
    if 'eom_ret' not in meta:
        meta['eom_ret']=meta.eom+pd.offsets.MonthEnd(1)
    else:meta['eom_ret']=pd.to_datetime(meta.eom_ret)
    if meta.duplicated(['id','eom']).any():raise ValueError('Duplicate stock-month keys')
    ids=meta.id.to_numpy(dtype=float)
    if not np.isfinite(ids).all() or not np.equal(ids,np.floor(ids)).all():
        raise ValueError('Invalid security identifiers')
    order=np.lexsort((ids,meta.eom.to_numpy()))
    raw=np.empty((len(meta),len(names)),dtype=np.float32)
    for start in range(0,len(names),50):
        end=min(start+50,len(names))
        raw[:,start:end]=chars[names[start:end]].to_numpy(dtype=np.float32)
    if not np.array_equal(order,np.arange(len(meta))):raw=raw[order]
    panel=pd.DataFrame(raw,columns=names,copy=False)
    for col in meta:panel[col]=meta.iloc[order][col].to_numpy()
    return panel,names,raw


def _prepare_daily(daily_ret,panel_ids):
    # Two bounded passes; no giant inverse-index arrays or private cache files.
    ids=np.sort(np.unique(panel_ids))
    dates=[]
    for start in range(0,len(daily_ret),250000):
        chunk=daily_ret.iloc[start:start+250000]
        d=pd.to_datetime(chunk['date']).to_numpy(dtype='datetime64[D]')
        dates.append(np.unique(d))
    if not dates:return np.empty((0,len(ids)),np.float32),np.array([],dtype='datetime64[D]'),ids
    dates=np.unique(np.concatenate(dates))
    matrix=np.full((len(dates),len(ids)),np.nan,dtype=np.float32)
    for start in range(0,len(daily_ret),250000):
        chunk=daily_ret.iloc[start:start+250000]
        rid=chunk['id'].to_numpy(dtype=float)
        columns=np.searchsorted(ids,rid)
        valid=columns<len(ids)
        valid[valid]=ids[columns[valid]]==rid[valid]
        d=pd.to_datetime(chunk['date']).to_numpy(dtype='datetime64[D]')
        returns=chunk['ret_exc'].to_numpy(dtype=float)
        matrix[np.searchsorted(dates,d[valid]),columns[valid]]=returns[valid].astype('float32')
    return matrix,dates,ids


def run_prepared(panel,names,raw,matrix,dates,daily_ids):
    np.random.seed(42)
    target=target_transform(panel)
    groups=panel.groupby('eom',sort=True).indices
    masks=test_flag(panel.ctff_test)
    has_test=bool(masks.any())
    coef_dates=[];coef_values=[]
    if SELECTED in ['ridge_daily','blend_daily']:
        _,_,importance=run_months(panel.groupby('eom',sort=True),names,diagnostics=True)
        coef_dates=[pd.Timestamp(v['eom']) for v in importance]
        coef_values=importance
    history={};outputs=[];tree=None;tree_year=None
    release_groups=panel.groupby('eom_ret',sort=True).indices
    released=set()
    for date,idx in groups.items():
        date=pd.Timestamp(date)
        for release,past_idx in release_groups.items():
            if release>date or release in released:continue
            released.add(release)
            past=panel.iloc[past_idx]
            valid=np.isfinite(past.ret_exc_lead1m.to_numpy(dtype=float))
            past=past.loc[valid]
            if not len(past):continue
            market=float(past.ret_exc_lead1m.mean())
            for stock,r in zip(past.id,past.ret_exc_lead1m):
                entries=history.setdefault(int(stock),deque(maxlen=60))
                entries.append((pd.Timestamp(release),float(r),market))
        selected=masks[idx] if has_test else np.ones(len(idx),dtype=bool)
        if not selected.any():continue
        frame=panel.iloc[idx]
        ids=frame.id.to_numpy(dtype=np.int64)
        ridge_signal=np.zeros(len(idx))
        if SELECTED in ['ridge_daily','blend_daily']:
            j=np.searchsorted(coef_dates,date,side='right')-1
            if j>=0:
                cf=coef_values[j]
                x=transform(frame,names)
                pfull=x@np.asarray(cf['coef_expanding'])
                precent=x@np.asarray(cf['coef_recent'])
                ridge_signal=standardized_signal(.5*(standardized_signal(pfull)+standardized_signal(precent)))
        tree_signal=np.zeros(len(idx))
        if SELECTED in ['tree_daily','tree_monthly','blend_daily']:
            if tree_year!=date.year:
                tree,_=fit_tree(raw,panel,target,date)
                tree_year=date.year
            if tree is not None:
                tree_signal=standardized_signal(tree.predict(np.asarray(raw[idx]),num_threads=4))
        signal=(.5*(ridge_signal+tree_signal) if SELECTED=='blend_daily' else
                ridge_signal if SELECTED=='ridge_daily' else tree_signal)
        if SELECTED=='tree_monthly':
            weights,_=risk_weights(ids,signal,history,date)
        else:
            weights,_=daily_covariance_weights(matrix,dates,daily_ids,ids,signal,date)
            if weights is None:weights,_=risk_weights(ids,signal,history,date)
        outputs.append(pd.DataFrame({'id':ids[selected],'eom':date,'w':weights[selected]}))
    if not outputs:raise ValueError('No formation observations')
    out=pd.concat(outputs,ignore_index=True)
    if out.isna().any().any() or not np.isfinite(out.w).all():raise ValueError('Invalid output weights')
    return out


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    if SELECTED=='original':return _original_main(chars,features,daily_ret)
    panel,names,raw=_prepare_monthly(chars,features)
    matrix,dates,daily_ids=_prepare_daily(daily_ret,panel.id.to_numpy(dtype=float))
    return run_prepared(panel,names,raw,matrix,dates,daily_ids)