"""Strategia 7: LSTM su features frattali multi-timeframe.
Usa sequenze di features frattali come input a un LSTM
per predire la direzione del prezzo.
"""
from __future__ import annotations
import sys
sys.path.insert(0, ".")

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
from src.data.downloader import load_data
from src.fractal.indicators import hurst_exponent, fractal_dimension_higuchi, volatility_ratio
from src.fractal.patterns import encode_candles, extract_body_ratios, extract_shadow_ratios

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")


class FractalLSTM(nn.Module):
    def __init__(self, input_size: int, hidden_size: int = 64, num_layers: int = 2, dropout: float = 0.3):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, (h_n, _) = self.lstm(x)
        out = self.classifier(h_n[-1])
        return out.squeeze(-1)


def extract_candle_features(df: pd.DataFrame, i: int) -> np.ndarray:
    """Extract per-candle features at index i."""
    o, h, l, c = df["open"].values[i], df["high"].values[i], df["low"].values[i], df["close"].values[i]
    v = df["volume"].values[i]
    total = h - l if h - l > 0 else 1e-10
    body = abs(c - o) / total
    upper_s = (h - max(o, c)) / total
    lower_s = (min(o, c) - l) / total
    direction = 1 if c > o else (-1 if c < o else 0)

    # Log return from previous candle
    if i > 0:
        prev_c = df["close"].values[i - 1]
        log_ret = np.log(c / prev_c) if prev_c > 0 else 0
    else:
        log_ret = 0

    return np.array([body, upper_s, lower_s, direction, log_ret, v])


def build_dataset(df: pd.DataFrame, seq_len: int = 48, lookahead: int = 6, min_ret: float = 0.003):
    """Build sequences of candle features with labels."""
    close = df["close"].values
    n = len(df)

    vol_mean = pd.Series(df["volume"].values).rolling(100, min_periods=1).mean().values

    sequences = []
    labels = []
    indices = []

    # Pre-compute additional features
    candle_types = encode_candles(df)
    body_ratios = extract_body_ratios(df)
    shadow_ratios = extract_shadow_ratios(df)

    for i in range(seq_len, n - lookahead, 2):
        seq = []
        for j in range(i - seq_len, i):
            feats = extract_candle_features(df, j)
            # Normalize volume by rolling mean
            feats[5] = feats[5] / vol_mean[j] if vol_mean[j] > 0 else 1.0
            seq.append(feats)

        future_ret = (close[i + lookahead - 1] - close[i - 1]) / close[i - 1]
        if abs(future_ret) < min_ret:
            continue

        sequences.append(seq)
        labels.append(1 if future_ret > 0 else 0)
        indices.append(i)

    return np.array(sequences), np.array(labels), np.array(indices)


print("=" * 60)
print("  STRATEGIA 7: LSTM FRACTAL — BTC 1H")
print("=" * 60)

df = load_data("BTC", "1h")
close = df["close"].values

SEQ_LEN = 48
LOOKAHEAD = 6
EPOCHS = 30
BATCH_SIZE = 256
LR = 0.001

print(f"\nSeq length: {SEQ_LEN}, Lookahead: {LOOKAHEAD}")
print("Building dataset...")

X, y, idx_arr = build_dataset(df, seq_len=SEQ_LEN, lookahead=LOOKAHEAD)
print(f"Samples: {len(X)}, Features per candle: {X.shape[2]}, Up ratio: {np.mean(y)*100:.1f}%")

# Chronological split
split = int(len(X) * 0.7)
val_split = int(len(X) * 0.85)

X_train, X_val, X_test = X[:split], X[split:val_split], X[val_split:]
y_train, y_val, y_test = y[:split], y[split:val_split], y[val_split:]
idx_test_arr = idx_arr[val_split:]

# Normalize features per-feature across time
n_features = X.shape[2]
for f in range(n_features):
    scaler = StandardScaler()
    X_train[:, :, f] = scaler.fit_transform(X_train[:, :, f])
    X_val[:, :, f] = scaler.transform(X_val[:, :, f])
    X_test[:, :, f] = scaler.transform(X_test[:, :, f])

# To tensors
X_train_t = torch.FloatTensor(X_train).to(DEVICE)
y_train_t = torch.FloatTensor(y_train).to(DEVICE)
X_val_t = torch.FloatTensor(X_val).to(DEVICE)
y_val_t = torch.FloatTensor(y_val).to(DEVICE)
X_test_t = torch.FloatTensor(X_test).to(DEVICE)

train_ds = TensorDataset(X_train_t, y_train_t)
train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)

# Model
model = FractalLSTM(input_size=n_features, hidden_size=64, num_layers=2, dropout=0.3).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-5)
criterion = nn.BCEWithLogitsLoss()
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)

print(f"\nTraining on {DEVICE}...")
best_val_acc = 0
patience_counter = 0

for epoch in range(EPOCHS):
    model.train()
    total_loss = 0
    for xb, yb in train_dl:
        optimizer.zero_grad()
        pred = model(xb)
        loss = criterion(pred, yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()

    # Validation
    model.eval()
    with torch.no_grad():
        val_pred = model(X_val_t)
        val_loss = criterion(val_pred, y_val_t).item()
        val_proba = torch.sigmoid(val_pred).cpu().numpy()
        val_acc = np.mean((val_proba > 0.5) == y_val)

    scheduler.step(val_loss)

    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), "data/processed/best_lstm.pt")
        patience_counter = 0
    else:
        patience_counter += 1

    if epoch % 5 == 0 or patience_counter > 8:
        print(f"  Epoch {epoch:2d}: train_loss={total_loss/len(train_dl):.4f} val_loss={val_loss:.4f} val_acc={val_acc*100:.1f}% best={best_val_acc*100:.1f}%")

    if patience_counter > 10:
        print(f"  Early stopping at epoch {epoch}")
        break

# Load best model and test
model.load_state_dict(torch.load("data/processed/best_lstm.pt", weights_only=True))
model.eval()
with torch.no_grad():
    test_pred = model(X_test_t)
    test_proba = torch.sigmoid(test_pred).cpu().numpy()

test_acc = np.mean((test_proba > 0.5) == y_test)
print(f"\nTest accuracy (base): {test_acc*100:.1f}%")

# Threshold sweep
print("\n--- THRESHOLD SWEEP ---")
for thr in [0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80]:
    accs = []
    capital = 1000
    n_trades = 0

    for j in range(len(X_test)):
        p = test_proba[j]
        i = idx_test_arr[j]
        actual = (close[i + LOOKAHEAD - 1] - close[i - 1]) / close[i - 1]

        if p >= thr:
            accs.append(1 if actual > 0 else 0)
            ret = actual - 0.002
            capital *= (1 + ret * 0.3)
            n_trades += 1
        elif p <= (1 - thr):
            accs.append(1 if actual < 0 else 0)
            ret = -actual - 0.002
            capital *= (1 + ret * 0.3)
            n_trades += 1

    if not accs:
        print(f"  thr={thr:.2f}: no signals")
        continue

    acc = np.mean(accs) * 100
    total_ret = (capital - 1000) / 1000 * 100
    # Annualized
    test_days = (idx_test_arr[-1] - idx_test_arr[0]) / 24
    years = test_days / 365.25 if test_days > 0 else 1
    ann_ret = ((capital / 1000) ** (1 / years) - 1) * 100 if years > 0 and capital > 0 else -100
    trades_yr = n_trades / years if years > 0 else 0
    print(f"  thr={thr:.2f}: trades={n_trades:5d} acc={acc:.1f}% ret={total_ret:+.1f}% ann={ann_ret:+.1f}% trades/yr={trades_yr:.0f}")

# Also try ETH
print("\n\n" + "=" * 60)
print("  LSTM SU ETH 1H (same model architecture)")
print("=" * 60)

df_eth = load_data("ETH", "1h")
close_eth = df_eth["close"].values

X_eth, y_eth, idx_eth = build_dataset(df_eth, seq_len=SEQ_LEN, lookahead=LOOKAHEAD)
print(f"ETH samples: {len(X_eth)}, Up ratio: {np.mean(y_eth)*100:.1f}%")

split_e = int(len(X_eth) * 0.7)
val_e = int(len(X_eth) * 0.85)
X_train_e, X_val_e, X_test_e = X_eth[:split_e], X_eth[split_e:val_e], X_eth[val_e:]
y_train_e, y_val_e, y_test_e = y_eth[:split_e], y_eth[split_e:val_e], y_eth[val_e:]
idx_test_e = idx_eth[val_e:]

for f in range(n_features):
    sc = StandardScaler()
    X_train_e[:, :, f] = sc.fit_transform(X_train_e[:, :, f])
    X_val_e[:, :, f] = sc.transform(X_val_e[:, :, f])
    X_test_e[:, :, f] = sc.transform(X_test_e[:, :, f])

X_tr_e = torch.FloatTensor(X_train_e).to(DEVICE)
y_tr_e = torch.FloatTensor(y_train_e).to(DEVICE)
X_va_e = torch.FloatTensor(X_val_e).to(DEVICE)
y_va_e = torch.FloatTensor(y_val_e).to(DEVICE)
X_te_e = torch.FloatTensor(X_test_e).to(DEVICE)

model_eth = FractalLSTM(input_size=n_features, hidden_size=64, num_layers=2, dropout=0.3).to(DEVICE)
opt_e = torch.optim.Adam(model_eth.parameters(), lr=LR, weight_decay=1e-5)
ds_e = TensorDataset(X_tr_e, y_tr_e)
dl_e = DataLoader(ds_e, batch_size=BATCH_SIZE, shuffle=True)
sch_e = torch.optim.lr_scheduler.ReduceLROnPlateau(opt_e, patience=5, factor=0.5)

best_e = 0
pc = 0
for epoch in range(EPOCHS):
    model_eth.train()
    tl = 0
    for xb, yb in dl_e:
        opt_e.zero_grad()
        p = model_eth(xb)
        loss = criterion(p, yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model_eth.parameters(), 1.0)
        opt_e.step()
        tl += loss.item()

    model_eth.eval()
    with torch.no_grad():
        vp = model_eth(X_va_e)
        vl = criterion(vp, y_va_e).item()
        va = np.mean((torch.sigmoid(vp).cpu().numpy() > 0.5) == y_val_e)

    sch_e.step(vl)
    if va > best_e:
        best_e = va
        torch.save(model_eth.state_dict(), "data/processed/best_lstm_eth.pt")
        pc = 0
    else:
        pc += 1
    if epoch % 5 == 0:
        print(f"  Epoch {epoch:2d}: val_acc={va*100:.1f}% best={best_e*100:.1f}%")
    if pc > 10:
        break

model_eth.load_state_dict(torch.load("data/processed/best_lstm_eth.pt", weights_only=True))
model_eth.eval()
with torch.no_grad():
    tp_e = torch.sigmoid(model_eth(X_te_e)).cpu().numpy()

print(f"\nETH Test accuracy: {np.mean((tp_e > 0.5) == y_test_e)*100:.1f}%")

for thr in [0.55, 0.60, 0.65, 0.70]:
    accs = []
    capital = 1000
    for j in range(len(X_test_e)):
        p = tp_e[j]
        i = idx_test_e[j]
        actual = (close_eth[i + LOOKAHEAD - 1] - close_eth[i - 1]) / close_eth[i - 1]
        if p >= thr:
            accs.append(1 if actual > 0 else 0)
            capital *= (1 + (actual - 0.002) * 0.3)
        elif p <= (1 - thr):
            accs.append(1 if actual < 0 else 0)
            capital *= (1 + (-actual - 0.002) * 0.3)
    if accs:
        print(f"  thr={thr:.2f}: trades={len(accs):5d} acc={np.mean(accs)*100:.1f}% ret={(capital-1000)/10:+.1f}%")