"""
Hybrid Gated Fusion Network for binary options strategy filtering.

Architecture: 3 input branches (strategy, context, instrument) →
  gate mechanism (context suppresses unreliable strategies) →
  shared fusion → 6 per-expiry sigmoid heads.

~18K parameters. Designed for the 28K-row dataset from build_dataset.py.
"""

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model, Input

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'


class FocalBinaryCrossentropy(tf.keras.losses.Loss):
    """Focal loss for extreme class imbalance (88% negative, 12% positive)."""
    def __init__(self, gamma=2.0, alpha=0.88, name='focal_bce', **kwargs):
        if 'reduction' in kwargs:
            kwargs.pop('reduction')
        super().__init__(name=name, **kwargs)
        self.gamma = gamma
        self.alpha = alpha

    def call(self, y_true, y_pred):
        y_pred = tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7)
        bce = -(y_true * tf.math.log(y_pred) + (1 - y_true) * tf.math.log(1 - y_pred))
        p_t = y_true * y_pred + (1 - y_true) * (1 - y_pred)
        focal_weight = tf.pow(1 - p_t, self.gamma)
        alpha_weight = y_true * self.alpha + (1 - y_true) * (1 - self.alpha)
        return alpha_weight * focal_weight * bce

    def get_config(self):
        config = super().get_config()
        config.update({'gamma': self.gamma, 'alpha': self.alpha})
        return config


def build_model(n_features_strategy=10, n_features_context=8, n_expiries=6, instrument_dim=2):
    """
    Hybrid Gated Fusion Network.

    Inputs:
      strategy_input  — 5×signal (-1/0/1) + 5×confidence (0-1)
      context_input   — RSI14, ATR14(scaled), MACD_hist(scaled), body_pct,
                         hour_sin, hour_cos, dow_sin, dow_cos
      instrument_input — one-hot [is_BITCOIN, is_EURUSD]

    Outputs: 6 × sigmoid → P(correct | expiry_Xh)
    """
    strategy_input = Input(shape=(n_features_strategy,), name='strategy_features')
    context_input = Input(shape=(n_features_context,), name='context_features')
    instrument_input = Input(shape=(instrument_dim,), name='instrument')

    # ── Strategy branch ──
    s = layers.Dense(64)(strategy_input)
    s = layers.BatchNormalization()(s)
    s = layers.ReLU()(s)
    s = layers.Dropout(0.3)(s)
    s = layers.Dense(32)(s)
    s = layers.BatchNormalization()(s)
    s = layers.ReLU()(s)
    s = layers.Dropout(0.3)(s)  # strategy_embedding (32)

    # ── Context branch ──
    c = layers.Dense(32)(context_input)
    c = layers.BatchNormalization()(c)
    c = layers.ReLU()(c)
    c = layers.Dropout(0.3)(c)
    c = layers.Dense(16)(c)
    c = layers.BatchNormalization()(c)
    c = layers.ReLU()(c)
    c = layers.Dropout(0.3)(c)  # context_embedding (16)

    # ── Gate: context controls trust in strategies ──
    gate_input = layers.Concatenate()([s, c, instrument_input])
    gate = layers.Dense(32, activation='sigmoid', name='gate')(gate_input)
    gated_strategy = layers.Multiply(name='gated_strategy')([s, gate])

    # ── Fusion ──
    fused = layers.Concatenate(name='fused')([gated_strategy, c, instrument_input])

    f = layers.Dense(64)(fused)
    f = layers.BatchNormalization()(f)
    f = layers.ReLU()(f)
    f = layers.Dropout(0.4)(f)

    f = layers.Dense(32)(f)
    f = layers.BatchNormalization()(f)
    f = layers.ReLU()(f)
    f = layers.Dropout(0.3)(f)  # shared_representation (32)

    # ── Per-expiry output heads ──
    expiry_names = ['1h', '2h', '3h', '4h', '5h', '10h']
    outputs = {}
    for exp in expiry_names:
        h = layers.Dense(16, name=f'exp_{exp}_hidden')(f)
        h = layers.ReLU()(h)
        h = layers.Dropout(0.2)(h)
        outputs[f'expiry_{exp}'] = layers.Dense(1, activation='sigmoid', name=f'expiry_{exp}')(h)

    model = Model(
        inputs=[strategy_input, context_input, instrument_input],
        outputs=outputs,
        name='strategy_filter'
    )

    return model


def compile_model(model, lr=1e-3, wd=1e-4):
    """Compile with Focal BCE per head + AdamW."""
    expiry_names = ['1h', '2h', '3h', '4h', '5h', '10h']
    losses = {f'expiry_{exp}': FocalBinaryCrossentropy(gamma=2.0, alpha=0.88) for exp in expiry_names}
    loss_weights = {f'expiry_{exp}': 1.0 for exp in expiry_names}

    optimizer = tf.keras.optimizers.AdamW(
        learning_rate=lr,
        weight_decay=wd,
        beta_1=0.9,
        beta_2=0.999,
    )

    model.compile(
        optimizer=optimizer,
        loss=losses,
        loss_weights=loss_weights,
        metrics={f'expiry_{exp}': ['binary_accuracy', tf.keras.metrics.AUC(name='auc')] for exp in expiry_names},
    )
    return model