"""Non-promoting, position-conditioned continuation-value experiment.

The value target compares selling now against a fixed short continuation, not
an oracle's best future exit. Inference always uses the candidate's own state.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd

from strategies.strategy_activation_scores import COMMISSION_RATE

HORIZON = 4
MAX_EXTENSION = 12
MIN_TRAIN_TRADES = 20
MARKET_NAMES = ('score', 'score_change', 'return_4', 'volatility_20',
                'mvrv_regime', 'rvol_norm', 'fed_net_liq_sign')
STATE_NAMES = ('age', 'unrealized_return', 'drawdown_from_high', 'extension_age',
               'age_times_return', 'return_squared', 'drawdown_squared')


def market_features(frame):
    columns = ['activation_score', 'close', 'mvrv_regime', 'rvol_norm', 'fed_net_liq_sign']
    for name in columns:
        if name not in frame:
            raise ValueError(f'Missing learned-exit input: {name}')
        if not np.isfinite(frame[name].to_numpy(dtype=float)).all():
            raise ValueError(f'Nonfinite learned-exit input: {name}')
    score = frame.activation_score / 1000
    returns = frame.close.pct_change()
    return np.column_stack([score, score.diff(2).fillna(0), frame.close.pct_change(4).fillna(0),
                            returns.rolling(20, min_periods=2).std().fillna(0),
                            frame.mvrv_regime, frame.rvol_norm, frame.fed_net_liq_sign])


def state_features(t, entry_bar, entry_price, high, close, pending):
    age = min((t - entry_bar) / 100, 5.)
    pnl = close / entry_price - 1
    dd = close / high - 1
    extension = 0. if pending < 0 else (t - pending) / MAX_EXTENSION
    return np.array([age, pnl, dd, extension, age*pnl, pnl*pnl, dd*dd])


def feature_vector(market, state, kind):
    if kind == 'market':
        return market
    if kind == 'state':
        return state
    if kind == 'combined':
        return np.r_[market, state, market[0]*state[1], market[4]*state[2]]
    raise ValueError(f'Unknown input kind: {kind}')


@dataclass
class ValueHead:
    kind: str
    mean: np.ndarray
    scale: np.ndarray
    coefficients: np.ndarray
    intercept: float
    tree: dict | None = None

    def predict(self, market, state):
        x = feature_vector(market, state, self.kind)
        z = (x-self.mean) / self.scale
        if self.tree is not None:
            # sklearn's tree predictor evaluates float32 features; preserve its
            # branch decisions when replaying the portable JSON representation.
            z = z.astype(np.float32)
            node = 0
            while self.tree['left'][node] != -1:
                branch = 'left' if z[self.tree['feature'][node]] <= self.tree['threshold'][node] else 'right'
                node = self.tree[branch][node]
            return self.tree['value'][node]
        return float(np.sum(z * self.coefficients) + self.intercept)

    def artifact(self):
        names = list(MARKET_NAMES if self.kind == 'market' else STATE_NAMES)
        if self.kind == 'combined':
            names = list(MARKET_NAMES + STATE_NAMES) + ['score_times_return', 'regime_times_drawdown']
        return {'kind': self.kind, 'features': names, 'mean': self.mean.tolist(),
                'scale': self.scale.tolist(), 'coefficients': self.coefficients.tolist(),
                'intercept': self.intercept, 'ridge_penalty': 1.0 if self.tree is None else None,
                'model_type': 'ridge' if self.tree is None else 'tree_depth_2', 'tree': self.tree}


def fit_head(records, times, cutoff, kind, model_type='ridge'):
    """Purge target endpoints at/after cutoff; equal total weight per trade."""
    eligible = [r for r in records if times[r['label_end']] < pd.Timestamp(cutoff)]
    groups = np.array([r['entry_bar'] for r in eligible], dtype=int)
    unique, counts = np.unique(groups, return_counts=True)
    audit = {'samples': len(eligible), 'trades': len(unique), 'cutoff': str(cutoff),
             'max_label_end': str(max((times[r['label_end']] for r in eligible), default=None)),
             'abstain': len(unique) < MIN_TRAIN_TRADES}
    if audit['abstain']:
        return None, audit
    count_by_trade = dict(zip(unique, counts))
    weights = np.array([1 / count_by_trade[g] for g in groups], dtype=float)
    weights /= weights.sum()
    x = np.array([feature_vector(r['market'], r['state'], kind) for r in eligible])
    y = np.array([r['target'] for r in eligible])
    if not np.isfinite(x).all() or not np.isfinite(y).all():
        raise ValueError('Nonfinite learned-exit training rows')
    mean = np.sum(weights[:, None]*x, axis=0)
    scale = np.sqrt(np.sum(weights[:, None]*((x-mean)**2), axis=0))
    scale[scale < 1e-8] = 1.
    z = (x-mean) / scale
    intercept = float(np.sum(weights*y))
    if model_type == 'tree':
        from sklearn.tree import DecisionTreeRegressor
        estimator = DecisionTreeRegressor(max_depth=2, min_weight_fraction_leaf=.1, random_state=0)
        estimator.fit(z, y, sample_weight=weights)
        tree = estimator.tree_
        exported = {'left': tree.children_left.tolist(), 'right': tree.children_right.tolist(),
                    'feature': tree.feature.tolist(), 'threshold': tree.threshold.tolist(),
                    'value': tree.value[:, 0, 0].tolist()}
        return ValueHead(kind, mean, scale, np.zeros(z.shape[1]), intercept, exported), audit
    if model_type != 'ridge':
        raise ValueError('Unknown value model type')
    gram = np.einsum('ni,n,nj->ij', z, weights, z)
    rhs = np.einsum('ni,n,n->i', z, weights, y-intercept)
    coefficients = np.linalg.solve(gram + np.eye(z.shape[1]), rhs)
    if not np.isfinite(coefficients).all():
        raise ValueError('Nonfinite learned-exit coefficients')
    return ValueHead(kind, mean, scale, coefficients, intercept), audit


def continuation_target(close, high, t, entry_bar, trade_high, trail, budget=MAX_EXTENSION):
    """Canonical after-fee wealth difference versus exit now; endpoint explicit."""
    end = min(t+HORIZON, t+budget)
    if budget < 1 or end >= len(close):
        return None
    wealth = 1.
    for j in range(t+1, end+1):
        if j != entry_bar+1:
            trade_high = max(trade_high, high[j])
        stop = trail > 0 and close[j] <= trade_high*(1-trail)
        last = stop or j == end
        # Fee is subtracted from the exit bar's return, matching calculate_metrics.
        wealth *= close[j]/close[j-1] - (COMMISSION_RATE if last else 0.)
        if last:
            return wealth - (1-COMMISSION_RATE), j
    raise AssertionError('Continuation must terminate')


def training_records(frame, opportunities):
    """Baseline on-policy rows plus bounded, stopped post-exit continuations."""
    market = market_features(frame)
    close, high = frame.close.to_numpy(), frame.high.to_numpy()
    records = []
    position, use_raw = False, False
    entry_bar, trade_high = -1, 0.

    def append(t, entry, peak, pending):
        budget = MAX_EXTENSION if pending < 0 else MAX_EXTENSION-(t-pending)
        target = continuation_target(close, high, t, entry, peak, opportunities['trail'], budget)
        if target is not None:
            records.append({'bar': t, 'entry_bar': entry, 'label_end': target[1], 'target': target[0],
                            'market': market[t], 'state': state_features(t, entry, close[entry], peak, close[t], pending),
                            'is_exit_event': pending == t})

    for t in range(len(frame)):
        if not position:
            if opportunities['entry'][t] and opportunities['time_ok'][t]:
                position, entry_bar, trade_high = True, t, close[t]
                use_raw = opportunities['dead_zone'] and opportunities['score_exit'][t] < opportunities['exit_threshold']
            continue
        if t != entry_bar+1:
            trade_high = max(trade_high, high[t])
        stop = opportunities['trail'] > 0 and close[t] <= trade_high*(1-opportunities['trail'])
        score_exit = opportunities['fallback'][t] if use_raw else opportunities['exit'][t]
        if not opportunities['time_ok'][t]:
            continue
        if stop:
            position = False
            continue
        append(t, entry_bar, trade_high, t if score_exit else -1)
        if score_exit:
            peak = trade_high
            for j in range(t+1, min(t+MAX_EXTENSION, len(frame))):
                if j != entry_bar+1:
                    peak = max(peak, high[j])
                if not opportunities['time_ok'][j] or (opportunities['trail'] > 0 and close[j] <= peak*(1-opportunities['trail'])):
                    break
                append(j, entry_bar, peak, t)
            position = False
    return records


def replay(frame, opportunities, *, start, end, head=None, mode='baseline'):
    """Sequential candidate rollout, flat at evaluation start; warm market data.

    `veto` evaluates at and after a score exit; `manage` also exits early.
    Missing fitted heads abstain to baseline. `delay` is the constant control.
    """
    if mode not in {'baseline', 'veto', 'manage', 'delay'}:
        raise ValueError('Invalid exit policy mode')
    result = frame.copy()
    market = market_features(frame)
    close, high = frame.close.to_numpy(), frame.high.to_numpy()
    allowed = ((frame.time >= pd.Timestamp(start)) &
               (frame.time < pd.Timestamp(end)+pd.Timedelta(days=1))).to_numpy() & opportunities['time_ok']
    n = len(frame)
    positions, entered, exited = np.zeros(n, int), np.zeros(n, bool), np.zeros(n, bool)
    decisions = []
    position, use_raw = False, False
    entry_bar, trade_high, pending = -1, 0., -1
    for t in range(n):
        if position:
            if t != entry_bar+1:
                trade_high = max(trade_high, high[t])
            stop = opportunities['trail'] > 0 and close[t] <= trade_high*(1-opportunities['trail'])
            score_exit = bool(opportunities['fallback'][t] if use_raw else opportunities['exit'][t])
            if allowed[t]:
                if score_exit and pending < 0:
                    pending = t
                state = state_features(t, entry_bar, close[entry_bar], trade_high, close[t], pending)
                value = None
                reason = ''
                if stop:
                    reason = 'trailing_stop'
                elif mode == 'baseline' or (head is None and mode != 'delay'):
                    if score_exit:
                        reason = 'score_exit'
                elif pending >= 0 and t-pending >= MAX_EXTENSION:
                    reason = 'extension_cap'
                elif mode != 'delay' and (pending >= 0 or mode == 'manage'):
                    value = head.predict(market[t], state)
                    if value <= 0:
                        reason = 'learned_exit'
                decisions.append({'bar': t, 'time': str(frame.time.iloc[t]), 'entry_bar': entry_bar,
                                  'age_bars': t-entry_bar, 'unrealized_return': state[1],
                                  'trade_high': trade_high,
                                  'score_exit': score_exit, 'pending_since': pending,
                                  'predicted_advantage': value, 'action': 'exit' if reason else 'hold',
                                  'reason': reason})
                if reason:
                    position, use_raw = False, False
                    exited[t], pending = True, -1
        elif allowed[t] and opportunities['entry'][t]:
            position, entry_bar, trade_high = True, t, close[t]
            entered[t] = True
            use_raw = opportunities['dead_zone'] and opportunities['score_exit'][t] < opportunities['exit_threshold']
        positions[t] = position
    result['position'], result['execute_entry'], result['execute_exit'] = positions, entered, exited
    return result, decisions