-- ================================================================================================
-- Migration 022: Ensemble ML Prediction Audit Tables
-- Production-ready schema for ensemble model attribution and A/B testing
-- Optimized with TimescaleDB hypertables for time-series queries
-- ================================================================================================

-- ================================================================================================
-- ENSEMBLE PREDICTIONS TABLE
-- Audit log for every ensemble prediction with per-model attribution
-- ================================================================================================
CREATE TABLE ensemble_predictions (
    -- Primary identifiers
    id UUID DEFAULT gen_random_uuid(),
    prediction_timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- Trading context
    symbol VARCHAR(20) NOT NULL,
    account_id VARCHAR(64),
    strategy_id VARCHAR(100),

    -- Ensemble decision
    ensemble_action VARCHAR(10) NOT NULL, -- BUY, SELL, HOLD
    ensemble_signal DOUBLE PRECISION NOT NULL CHECK (ensemble_signal >= -1.0 AND ensemble_signal <= 1.0),
    ensemble_confidence DOUBLE PRECISION NOT NULL CHECK (ensemble_confidence >= 0.0 AND ensemble_confidence <= 1.0),
    disagreement_rate DOUBLE PRECISION NOT NULL CHECK (disagreement_rate >= 0.0 AND disagreement_rate <= 1.0),

    -- Per-model votes (DQN)
    dqn_signal DOUBLE PRECISION,
    dqn_confidence DOUBLE PRECISION,
    dqn_weight DOUBLE PRECISION,
    dqn_vote VARCHAR(10), -- BUY, SELL, HOLD

    -- Per-model votes (PPO)
    ppo_signal DOUBLE PRECISION,
    ppo_confidence DOUBLE PRECISION,
    ppo_weight DOUBLE PRECISION,
    ppo_vote VARCHAR(10),

    -- Per-model votes (MAMBA-2)
    mamba2_signal DOUBLE PRECISION,
    mamba2_confidence DOUBLE PRECISION,
    mamba2_weight DOUBLE PRECISION,
    mamba2_vote VARCHAR(10),

    -- Per-model votes (TFT)
    tft_signal DOUBLE PRECISION,
    tft_confidence DOUBLE PRECISION,
    tft_weight DOUBLE PRECISION,
    tft_vote VARCHAR(10),

    -- Execution tracking (link to actual trades)
    order_id UUID REFERENCES orders(id) ON DELETE SET NULL,
    executed_price BIGINT, -- In cents or smallest unit
    position_size BIGINT, -- Quantity traded
    pnl BIGINT, -- Profit/Loss in cents (populated after trade closes)
    commission BIGINT DEFAULT 0, -- Transaction costs in cents
    slippage_bps INTEGER, -- Slippage in basis points (100 bps = 1%)

    -- A/B testing metadata
    ab_test_id UUID,
    ab_group VARCHAR(20), -- control, treatment
    ab_variant VARCHAR(50), -- Additional variant identifier

    -- Feature snapshot (for reproducibility and debugging)
    feature_snapshot JSONB, -- All input features used for this prediction

    -- Model checkpoint information
    dqn_checkpoint_id VARCHAR(255),
    ppo_checkpoint_id VARCHAR(255),
    mamba2_checkpoint_id VARCHAR(255),
    tft_checkpoint_id VARCHAR(255),

    -- System context
    node_id VARCHAR(50), -- Which server generated this prediction
    inference_latency_us INTEGER, -- Total inference time in microseconds
    aggregation_latency_us INTEGER, -- Time to aggregate model votes

    -- Compliance and audit
    user_id VARCHAR(64),
    session_id UUID,
    request_id UUID,

    -- Metadata
    metadata JSONB, -- Additional flexible metadata

    CONSTRAINT chk_ensemble_action CHECK (ensemble_action IN ('BUY', 'SELL', 'HOLD')),
    CONSTRAINT chk_model_votes CHECK (
        dqn_vote IS NULL OR dqn_vote IN ('BUY', 'SELL', 'HOLD')
    ),
    CONSTRAINT chk_valid_latency CHECK (
        inference_latency_us IS NULL OR inference_latency_us > 0
    )
,
    PRIMARY KEY (id, prediction_timestamp)
);

-- Indexes for fast queries
CREATE INDEX idx_ensemble_predictions_timestamp ON ensemble_predictions (prediction_timestamp DESC);
CREATE INDEX idx_ensemble_predictions_symbol_timestamp ON ensemble_predictions (symbol, prediction_timestamp DESC);
CREATE INDEX idx_ensemble_predictions_order_id ON ensemble_predictions (order_id) WHERE order_id IS NOT NULL;
CREATE INDEX idx_ensemble_predictions_ab_test ON ensemble_predictions (ab_test_id, ab_group) WHERE ab_test_id IS NOT NULL;
CREATE INDEX idx_ensemble_predictions_action ON ensemble_predictions (ensemble_action);
CREATE INDEX idx_ensemble_predictions_high_disagreement ON ensemble_predictions (disagreement_rate DESC) WHERE disagreement_rate > 0.5;

-- GIN index for JSONB feature snapshot queries
CREATE INDEX idx_ensemble_predictions_feature_snapshot ON ensemble_predictions USING GIN (feature_snapshot);

-- Index for P&L attribution queries
CREATE INDEX idx_ensemble_predictions_pnl ON ensemble_predictions (pnl DESC NULLS LAST) WHERE pnl IS NOT NULL;

-- TimescaleDB hypertable for time-series optimization
SELECT create_hypertable('ensemble_predictions', 'prediction_timestamp',
    chunk_time_interval => INTERVAL '1 day',
    if_not_exists => TRUE
);

-- Compress old data (older than 7 days) to save space

COMMENT ON TABLE ensemble_predictions IS 'Audit log of every ensemble prediction with per-model attribution and execution tracking';
COMMENT ON COLUMN ensemble_predictions.ensemble_signal IS 'Weighted average signal from -1.0 (strong sell) to 1.0 (strong buy)';
COMMENT ON COLUMN ensemble_predictions.disagreement_rate IS 'Percentage of models that disagree with ensemble decision (0.0-1.0)';
COMMENT ON COLUMN ensemble_predictions.feature_snapshot IS 'JSONB snapshot of all input features for reproducibility';
COMMENT ON COLUMN ensemble_predictions.inference_latency_us IS 'Total time for all models to generate predictions (microseconds)';

-- ================================================================================================
-- MODEL PERFORMANCE ATTRIBUTION TABLE
-- Rolling performance metrics per model for adaptive weighting
-- ================================================================================================
CREATE TABLE model_performance_attribution (
    -- Primary identifiers
    id UUID DEFAULT gen_random_uuid(),
    prediction_timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- Model identification
    model_id VARCHAR(50) NOT NULL, -- DQN, PPO, MAMBA2, TFT
    symbol VARCHAR(20) NOT NULL,

    -- Performance metrics (rolling window)
    window_hours INTEGER NOT NULL, -- 1, 24, 168 (1h, 1d, 1w)
    total_predictions INTEGER NOT NULL DEFAULT 0,
    correct_predictions INTEGER NOT NULL DEFAULT 0,
    accuracy DOUBLE PRECISION NOT NULL DEFAULT 0.0,

    -- P&L metrics
    total_pnl BIGINT NOT NULL DEFAULT 0, -- Total P&L in cents
    total_return DOUBLE PRECISION DEFAULT 0.0, -- Total return percentage
    sharpe_ratio DOUBLE PRECISION, -- Risk-adjusted returns
    sortino_ratio DOUBLE PRECISION, -- Downside risk-adjusted returns
    max_drawdown DOUBLE PRECISION, -- Maximum peak-to-trough decline
    win_rate DOUBLE PRECISION, -- Percentage of profitable trades

    -- Trading metrics
    avg_trade_pnl BIGINT, -- Average P&L per trade in cents
    total_trades INTEGER DEFAULT 0,
    winning_trades INTEGER DEFAULT 0,
    losing_trades INTEGER DEFAULT 0,

    -- Contribution to ensemble
    avg_weight DOUBLE PRECISION NOT NULL DEFAULT 0.0, -- Average weight in ensemble
    avg_confidence DOUBLE PRECISION NOT NULL DEFAULT 0.0, -- Average prediction confidence
    avg_signal DOUBLE PRECISION, -- Average signal strength

    -- Disagreement patterns
    disagreement_count INTEGER DEFAULT 0, -- Times model disagreed with ensemble
    disagreement_rate DOUBLE PRECISION, -- Percentage of disagreements

    -- Model checkpoint
    checkpoint_id VARCHAR(255), -- Active checkpoint during this window

    -- Metadata
    metadata JSONB,

    -- Constraints
    CONSTRAINT chk_model_id CHECK (model_id IN ('DQN', 'PPO', 'MAMBA2', 'TFT')),
    CONSTRAINT chk_window_hours CHECK (window_hours IN (1, 24, 168)),
    CONSTRAINT chk_accuracy_range CHECK (accuracy >= 0.0 AND accuracy <= 1.0),
    CONSTRAINT chk_prediction_counts CHECK (correct_predictions <= total_predictions)
,
    PRIMARY KEY (id, prediction_timestamp)
);

-- Indexes for fast queries
CREATE INDEX idx_model_performance_model_timestamp ON model_performance_attribution (model_id, prediction_timestamp DESC);
CREATE INDEX idx_model_performance_symbol_timestamp ON model_performance_attribution (symbol, prediction_timestamp DESC);
CREATE INDEX idx_model_performance_window ON model_performance_attribution (window_hours, prediction_timestamp DESC);
CREATE INDEX idx_model_performance_sharpe ON model_performance_attribution (sharpe_ratio DESC NULLS LAST);
CREATE INDEX idx_model_performance_accuracy ON model_performance_attribution (accuracy DESC);

-- Composite index for model comparison queries
CREATE INDEX idx_model_performance_comparison ON model_performance_attribution (symbol, window_hours, prediction_timestamp DESC);

-- TimescaleDB hypertable for time-series optimization
SELECT create_hypertable('model_performance_attribution', 'prediction_timestamp',
    chunk_time_interval => INTERVAL '1 day',
    if_not_exists => TRUE
);

-- Compress old data (older than 30 days)

COMMENT ON TABLE model_performance_attribution IS 'Rolling performance metrics per model for adaptive weight adjustment';
COMMENT ON COLUMN model_performance_attribution.window_hours IS 'Rolling window size: 1h, 24h, or 168h (1 week)';
COMMENT ON COLUMN model_performance_attribution.sharpe_ratio IS 'Annualized risk-adjusted returns (mean return / std dev of returns)';
COMMENT ON COLUMN model_performance_attribution.avg_weight IS 'Average weight assigned to this model in ensemble decisions';

-- ================================================================================================
-- A/B TEST EXPERIMENTS TABLE
-- Track A/B test configurations and status
-- ================================================================================================
CREATE TABLE ab_test_experiments (
    -- Primary identifiers
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    test_id UUID NOT NULL,

    -- Test configuration
    test_name VARCHAR(200) NOT NULL,
    description TEXT,
    control_variant VARCHAR(50) NOT NULL, -- e.g., "DQN_ONLY"
    treatment_variant VARCHAR(50) NOT NULL, -- e.g., "ENSEMBLE"

    -- Test parameters
    traffic_split DOUBLE PRECISION NOT NULL DEFAULT 0.5 CHECK (traffic_split >= 0.0 AND traffic_split <= 1.0),
    min_sample_size INTEGER NOT NULL DEFAULT 1000,
    significance_level DOUBLE PRECISION NOT NULL DEFAULT 0.05,
    max_duration_hours INTEGER NOT NULL DEFAULT 168, -- 1 week default

    -- Status tracking
    status VARCHAR(20) NOT NULL DEFAULT 'draft', -- draft, running, paused, completed, cancelled
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,

    -- Results (populated when test completes)
    control_predictions INTEGER DEFAULT 0,
    treatment_predictions INTEGER DEFAULT 0,
    control_sharpe DOUBLE PRECISION,
    treatment_sharpe DOUBLE PRECISION,
    sharpe_lift DOUBLE PRECISION, -- (treatment - control) / control
    pvalue DOUBLE PRECISION, -- Statistical significance
    is_significant BOOLEAN,
    recommendation TEXT, -- Human-readable recommendation

    -- Metadata
    created_by VARCHAR(64),
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    metadata JSONB,

    CONSTRAINT chk_ab_status CHECK (status IN ('draft', 'running', 'paused', 'completed', 'cancelled')),
    CONSTRAINT chk_variant_different CHECK (control_variant != treatment_variant)
);

CREATE INDEX idx_ab_test_experiments_test_id ON ab_test_experiments (test_id);
CREATE INDEX idx_ab_test_experiments_status ON ab_test_experiments (status) WHERE status = 'running';
CREATE INDEX idx_ab_test_experiments_started_at ON ab_test_experiments (started_at DESC);

COMMENT ON TABLE ab_test_experiments IS 'A/B test experiment configurations and results';
COMMENT ON COLUMN ab_test_experiments.traffic_split IS 'Percentage of traffic allocated to treatment (0.0-1.0)';
COMMENT ON COLUMN ab_test_experiments.sharpe_lift IS 'Percentage improvement: (treatment - control) / control';

-- ================================================================================================
-- CONTINUOUS AGGREGATES (TimescaleDB)
-- Pre-computed views for fast dashboard queries
-- ================================================================================================


-- ================================================================================================
-- UTILITY FUNCTIONS
-- Helper functions for common queries
-- ================================================================================================

-- Function: Get top performing models in last 24 hours
DROP FUNCTION IF EXISTS get_top_models_24h(VARCHAR, INTEGER) CASCADE;
CREATE FUNCTION get_top_models_24h(
    p_symbol VARCHAR(20) DEFAULT NULL,
    p_limit INTEGER DEFAULT 5
)
RETURNS TABLE (
    model_id VARCHAR(50),
    total_predictions INTEGER,
    accuracy DOUBLE PRECISION,
    sharpe_ratio DOUBLE PRECISION,
    total_pnl BIGINT,
    avg_weight DOUBLE PRECISION
) AS $$
BEGIN
    RETURN QUERY
    SELECT
        mpa.model_id,
        mpa.total_predictions,
        mpa.accuracy,
        mpa.sharpe_ratio,
        mpa.total_pnl,
        mpa.avg_weight
    FROM model_performance_attribution mpa
    WHERE
        mpa.window_hours = 24
        AND mpa.prediction_timestamp >= NOW() - INTERVAL '24 hours'
        AND (p_symbol IS NULL OR mpa.symbol = p_symbol)
    ORDER BY mpa.sharpe_ratio DESC NULLS LAST
    LIMIT p_limit;
END;
$$ LANGUAGE plpgsql;

COMMENT ON FUNCTION get_top_models_24h IS 'Get top N performing models in last 24 hours by Sharpe ratio';

-- Function: Calculate model correlation matrix (last 7 days)
DROP FUNCTION IF EXISTS calculate_model_correlation_7d(VARCHAR) CASCADE;
CREATE FUNCTION calculate_model_correlation_7d(
    p_symbol VARCHAR(20) DEFAULT NULL
)
RETURNS TABLE (
    model_a VARCHAR(50),
    model_b VARCHAR(50),
    correlation DOUBLE PRECISION,
    sample_size INTEGER
) AS $$
BEGIN
    -- This is a simplified version; full Pearson correlation would require more complex SQL
    -- For production, consider computing this in application code or using PostgreSQL extensions
    RETURN QUERY
    WITH model_signals AS (
        SELECT
            prediction_timestamp,
            symbol,
            dqn_signal,
            ppo_signal,
            mamba2_signal,
            tft_signal
        FROM ensemble_predictions
        WHERE
            timestamp >= NOW() - INTERVAL '7 days'
            AND (p_symbol IS NULL OR symbol = p_symbol)
            AND dqn_signal IS NOT NULL
            AND ppo_signal IS NOT NULL
            AND mamba2_signal IS NOT NULL
            AND tft_signal IS NOT NULL
    )
    SELECT
        'DQN' AS model_a,
        'PPO' AS model_b,
        CORR(dqn_signal, ppo_signal) AS correlation,
        COUNT(*)::INTEGER AS sample_size
    FROM model_signals
    UNION ALL
    SELECT 'DQN', 'MAMBA2', CORR(dqn_signal, mamba2_signal), COUNT(*)::INTEGER FROM model_signals
    UNION ALL
    SELECT 'DQN', 'TFT', CORR(dqn_signal, tft_signal), COUNT(*)::INTEGER FROM model_signals
    UNION ALL
    SELECT 'PPO', 'MAMBA2', CORR(ppo_signal, mamba2_signal), COUNT(*)::INTEGER FROM model_signals
    UNION ALL
    SELECT 'PPO', 'TFT', CORR(ppo_signal, tft_signal), COUNT(*)::INTEGER FROM model_signals
    UNION ALL
    SELECT 'MAMBA2', 'TFT', CORR(mamba2_signal, tft_signal), COUNT(*)::INTEGER FROM model_signals;
END;
$$ LANGUAGE plpgsql;

COMMENT ON FUNCTION calculate_model_correlation_7d IS 'Calculate pairwise correlation between model signals (last 7 days)';

-- Function: Get high disagreement events (last 24 hours)
DROP FUNCTION IF EXISTS get_high_disagreement_events_24h(VARCHAR, DOUBLE PRECISION, INTEGER) CASCADE;
CREATE FUNCTION get_high_disagreement_events_24h(
    p_symbol VARCHAR(20) DEFAULT NULL,
    p_disagreement_threshold DOUBLE PRECISION DEFAULT 0.5,
    p_limit INTEGER DEFAULT 100
)
RETURNS TABLE (
    prediction_timestamp TIMESTAMPTZ,
    symbol VARCHAR(20),
    ensemble_action VARCHAR(10),
    ensemble_confidence DOUBLE PRECISION,
    disagreement_rate DOUBLE PRECISION,
    dqn_vote VARCHAR(10),
    ppo_vote VARCHAR(10),
    mamba2_vote VARCHAR(10),
    tft_vote VARCHAR(10)
) AS $$
BEGIN
    RETURN QUERY
    SELECT
        ep.prediction_timestamp,
        ep.symbol,
        ep.ensemble_action,
        ep.ensemble_confidence,
        ep.disagreement_rate,
        ep.dqn_vote,
        ep.ppo_vote,
        ep.mamba2_vote,
        ep.tft_vote
    FROM ensemble_predictions ep
    WHERE
        ep.prediction_timestamp >= NOW() - INTERVAL '24 hours'
        AND (p_symbol IS NULL OR ep.symbol = p_symbol)
        AND ep.disagreement_rate >= p_disagreement_threshold
    ORDER BY ep.disagreement_rate DESC, ep.prediction_timestamp DESC
    LIMIT p_limit;
END;
$$ LANGUAGE plpgsql;

COMMENT ON FUNCTION get_high_disagreement_events_24h IS 'Get predictions with high model disagreement (possible regime shifts)';

-- ================================================================================================
-- PERMISSIONS (Adjust based on your security model)
-- ================================================================================================

-- Grant read access to trading service
GRANT SELECT ON ensemble_predictions TO foxhunt;
GRANT SELECT ON model_performance_attribution TO foxhunt;
GRANT SELECT ON ab_test_experiments TO foxhunt;

-- Grant write access for predictions and performance updates
GRANT INSERT, UPDATE ON ensemble_predictions TO foxhunt;
GRANT INSERT, UPDATE ON model_performance_attribution TO foxhunt;
GRANT INSERT, UPDATE ON ab_test_experiments TO foxhunt;

-- Grant execute permissions on utility functions
GRANT EXECUTE ON FUNCTION get_top_models_24h TO foxhunt;
GRANT EXECUTE ON FUNCTION calculate_model_correlation_7d TO foxhunt;
GRANT EXECUTE ON FUNCTION get_high_disagreement_events_24h TO foxhunt;

-- ================================================================================================
-- END MIGRATION 022
-- ================================================================================================