# =========================
# 0. REPRODUCIBILITY
# =========================
import torch, numpy as np, random
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

# =========================
# 1. SAFE DOWNLOAD FUNCTION (FIXED)
# =========================
import os, requests, zipfile, time

def safe_download(url, path):
    if os.path.exists(path):
        print(f"{path} already exists, skipping download.")
        return

    for i in range(5):
        try:
            print(f"Downloading (attempt {i+1})...")
            r = requests.get(url, stream=True, timeout=60)

            with open(path, "wb") as f:
                for chunk in r.iter_content(8192):
                    if chunk:
                        f.write(chunk)

            print("Download complete")
            return

        except Exception as e:
            print("Retrying due to:", e)
            time.sleep(5)

    raise Exception("Download failed after retries")

# =========================
# 2. LOAD LABELS
# =========================
import pandas as pd

def load_split(csv_file):
    df = pd.read_csv(csv_file)
    df.columns = df.columns.str.strip()

    split = {}
    for _, row in df.iterrows():
        try:
            pid = str(int(row["Participant_ID"]))
            split[pid] = int(row["PHQ8_Binary"])
        except:
            continue
    return split

train_labels = load_split("train_split_Depression_AVEC2017.csv")
dev_labels = load_split("dev_split_Depression_AVEC2017.csv")

ALL_REQUIRED_IDS = set(list(train_labels.keys()) + list(dev_labels.keys()))
print("Total required participants:", len(ALL_REQUIRED_IDS))

# =========================
# 3. SELECTIVE EXTRACTION
# =========================
def extract_needed(zip_path):
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        for file in zip_ref.namelist():

            parts = file.split("/")
            if len(parts) < 2:
                continue

            folder = parts[1] if parts[0].startswith("DAIC") else parts[0]

            if any(folder.startswith(pid + "_") for pid in ALL_REQUIRED_IDS):
                zip_ref.extract(file, "data")

# =========================
# 4. DOWNLOAD + EXTRACT (SAFE)
# =========================
urls = [
    "https://huggingface.co/datasets/ananyakarn/DAIC_WOZ_Data/resolve/main/DAIC_WOZ_Data.zip",
    "https://huggingface.co/datasets/ananyakarn/DAIC_WOZ_Data/resolve/main/DAIC_WOZ_Data-2.zip"
]

os.makedirs("data", exist_ok=True)

# 👉 skip extraction if already done
if len(os.listdir("data")) == 0:

    for i, url in enumerate(urls):
        zip_path = f"temp_{i}.zip"

        safe_download(url, zip_path)

        print(f"Extracting dataset {i+1}...")
        extract_needed(zip_path)

        os.remove(zip_path)

else:
    print("Dataset already extracted. Skipping download.")

# =========================
# 5. GET PARTICIPANTS
# =========================
def get_all_paths():
    paths = []
    for root, dirs, _ in os.walk("data"):
        for d in dirs:
            if "_P" in d or "_C" in d:
                paths.append(os.path.join(root, d))
    return paths

ALL_PATHS = get_all_paths()
print("Extracted participants:", len(ALL_PATHS))

# =========================
# 6. LIBRARIES
# =========================
import librosa
import torch.nn as nn
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics import accuracy_score, f1_score
from sklearn.preprocessing import StandardScaler

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# =========================
# 7. TEXT MODEL
# =========================
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
bert = AutoModel.from_pretrained("distilbert-base-uncased").to(device)
bert.eval()

# =========================
# 8. FEATURE FUNCTIONS
# =========================
def load_text(folder):
    try:
        file = [f for f in os.listdir(folder) if "TRANSCRIPT" in f][0]
        df = pd.read_csv(os.path.join(folder, file))
        return " ".join(df.iloc[:, -1].astype(str))
    except:
        return ""

def get_text_embedding(text):
    if text == "":
        return np.zeros(768)

    inputs = tokenizer(text, return_tensors="pt",
                       truncation=True, padding=True, max_length=256).to(device)
    with torch.no_grad():
        out = bert(**inputs)
    return out.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()

def get_audio_features(folder):
    try:
        file = [f for f in os.listdir(folder) if f.endswith("_AUDIO.wav")][0]
        y, sr = librosa.load(os.path.join(folder, file), sr=16000)
        mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=40)
        return np.mean(mfcc.T, axis=0)
    except:
        return np.zeros(40)

def get_visual_features(folder):
    feats = []

    for key in ["AUs", "pose", "gaze"]:
        try:
            file = [f for f in os.listdir(folder) if key in f][0]
            df = pd.read_csv(os.path.join(folder, file))
            df = df.select_dtypes(include=[np.number])

            df.replace(-100, np.nan, inplace=True)
            df.fillna(0, inplace=True)

            if df.shape[1] == 0:
                feats.append(np.zeros(20))
                continue

            feats.append(np.concatenate([df.mean().values, df.std().values]))

        except:
            feats.append(np.zeros(20))

    return np.concatenate(feats)

# =========================
# 9. BUILD DATA
# =========================
def build(label_dict):
    Xt, Xa, Xv, y = [], [], [], []

    for path in tqdm(ALL_PATHS):
        pid = os.path.basename(path).split("_")[0]

        if pid not in label_dict:
            continue

        Xt.append(get_text_embedding(load_text(path)))
        Xa.append(get_audio_features(path))
        Xv.append(get_visual_features(path))
        y.append(label_dict[pid])

    return np.array(Xt), np.array(Xa), np.array(Xv), np.array(y)

Xt, Xa, Xv, y_train = build(train_labels)
Xt_test, Xa_test, Xv_test, y_test = build(dev_labels)

if len(y_train) == 0:
    raise Exception("No training data found!")

print("Train size:", len(y_train))
print("Test size:", len(y_test))

# =========================
# 10. NORMALIZE
# =========================
sc_t, sc_a, sc_v = StandardScaler(), StandardScaler(), StandardScaler()

Xt = sc_t.fit_transform(Xt)
Xa = sc_a.fit_transform(Xa)
Xv = sc_v.fit_transform(Xv)

Xt_test = sc_t.transform(Xt_test)
Xa_test = sc_a.transform(Xa_test)
Xv_test = sc_v.transform(Xv_test)

# =========================
# 🔥 RANDOM FOREST (SAFE ADD)
# =========================
from sklearn.ensemble import RandomForestClassifier

# IMPORTANT: use numpy arrays BEFORE tensor conversion
# (at this point Xt, Xa, Xv are still numpy)

# Optional: reduce dimensionality for stability
Xt_rf = Xt[:, :128]     # reduce text features
Xv_rf = Xv[:, :50]      # reduce visual features

Xt_test_rf = Xt_test[:, :128]
Xv_test_rf = Xv_test[:, :50]

# Combine features
X_train_rf = np.concatenate([Xt_rf, Xa, Xv_rf], axis=1)
X_test_rf = np.concatenate([Xt_test_rf, Xa_test, Xv_test_rf], axis=1)

# Model
rf = RandomForestClassifier(
    n_estimators=200,
    max_depth=5,
    random_state=42
)

print("\nTraining Random Forest...")
rf.fit(X_train_rf, y_train)

# Predictions
rf_preds = rf.predict(X_test_rf)

rf_acc = accuracy_score(y_test, rf_preds)
rf_f1 = f1_score(y_test, rf_preds)

# =========================
# 11. MODELS
# =========================
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, vdim):
        super().__init__()
        self.t = nn.Sequential(nn.Linear(768,128), nn.ReLU())
        self.a = nn.Sequential(nn.Linear(40,32), nn.ReLU())
        self.v = nn.Sequential(nn.Linear(vdim,64), nn.ReLU())

        self.f = nn.Sequential(
            nn.Linear(224,64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64,1)
        )

    def forward(self, t, a, v):
        return self.f(torch.cat([self.t(t), self.a(a), self.v(v)], dim=1))


class AttentionFusionModel(nn.Module):
    def __init__(self, vdim):
        super().__init__()

        self.t = nn.Sequential(nn.Linear(768,128), nn.ReLU())
        self.a = nn.Sequential(nn.Linear(40,32), nn.ReLU())
        self.v = nn.Sequential(nn.Linear(vdim,64), nn.ReLU())

        self.attn = nn.Sequential(
            nn.Linear(224,64),
            nn.Tanh(),
            nn.Linear(64,3)
        )

        self.f = nn.Sequential(
            nn.Linear(224,64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64,1)
        )

    def forward(self, t, a, v):
        t_feat = self.t(t)
        a_feat = self.a(a)
        v_feat = self.v(v)

        combined = torch.cat([t_feat, a_feat, v_feat], dim=1)

        weights = torch.softmax(self.attn(combined), dim=1)

        fused = torch.cat([
            weights[:,0:1] * t_feat,
            weights[:,1:2] * a_feat,
            weights[:,2:3] * v_feat
        ], dim=1)

        return self.f(fused)


# =========================
# 12. CONVERT TO TENSORS
# =========================
Xt = torch.tensor(Xt, dtype=torch.float32).to(device)
Xa = torch.tensor(Xa, dtype=torch.float32).to(device)
Xv = torch.tensor(Xv, dtype=torch.float32).to(device)
yt = torch.tensor(y_train, dtype=torch.float32).to(device)

Xt_test = torch.tensor(Xt_test, dtype=torch.float32).to(device)
Xa_test = torch.tensor(Xa_test, dtype=torch.float32).to(device)
Xv_test = torch.tensor(Xv_test, dtype=torch.float32).to(device)


# =========================
# 13. TRAIN BASELINE MODEL
# =========================
baseline_model = Model(Xv.shape[1]).to(device)

opt1 = torch.optim.Adam(baseline_model.parameters(), lr=1e-4)
loss_fn = nn.BCEWithLogitsLoss()

print("\nTraining Baseline Model...")

for e in range(5):   # 🔥 reduced epochs (important)
    baseline_model.train()
    opt1.zero_grad()

    outputs = baseline_model(Xt, Xa, Xv).squeeze()
    loss = loss_fn(outputs, yt)

    loss.backward()
    opt1.step()

    print(f"Epoch {e+1}, Loss: {loss.item():.4f}")


# =========================
# 14. TRAIN ATTENTION MODEL
# =========================
attention_model = AttentionFusionModel(Xv.shape[1]).to(device)

opt2 = torch.optim.AdamW(attention_model.parameters(), lr=1e-4)
loss_fn_attn = nn.BCEWithLogitsLoss()

print("\nTraining Attention Model...")

for e in range(5):   # 🔥 reduced epochs
    attention_model.train()
    opt2.zero_grad()

    outputs = attention_model(Xt, Xa, Xv).squeeze()
    loss = loss_fn_attn(outputs, yt)

    loss.backward()
    opt2.step()

    print(f"Epoch {e+1}, Loss: {loss.item():.4f}")


# =========================
# 15. EVALUATION
# =========================
from sklearn.metrics import accuracy_score, f1_score

baseline_model.eval()
attention_model.eval()

with torch.no_grad():
    pred_baseline = (torch.sigmoid(
        baseline_model(Xt_test, Xa_test, Xv_test).squeeze()
    ) > 0.5).int().cpu().numpy()

    pred_attention = (torch.sigmoid(
        attention_model(Xt_test, Xa_test, Xv_test).squeeze()
    ) > 0.5).int().cpu().numpy()


print("\n===== MODEL COMPARISON =====")

print("\nRandom Forest:")
print("Accuracy:", rf_acc)
print("F1 Score:", rf_f1)

print("\nBaseline Model:")
print("Accuracy:", accuracy_score(y_test, pred_baseline))
print("F1 Score:", f1_score(y_test, pred_baseline))

print("\nAttention Model:")
print("Accuracy:", accuracy_score(y_test, pred_attention))
print("F1 Score:", f1_score(y_test, pred_attention))