"""
MindCheck Flask Backend — serves PHQ-8 predictions from the trained
MPNetMultiOutputRegressor model.

Endpoints:
    POST /predict   — accepts {"answers": ["...", ...]} (8 text answers)
    GET  /health    — health check
"""

import os
import re
import numpy as np
import torch
import torch.nn as nn
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModel

# ═══════════════════════════════════════════════════════════════════════════════
# 1. CONFIG
# ═══════════════════════════════════════════════════════════════════════════════

MODEL_PATH = os.environ.get(
    "MODEL_PATH",
    os.path.join(os.path.dirname(__file__), "..", "NOTEBOOK&MODEL", "model.pt"),
)
MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
MAX_LEN = 512
HEAD_TOKENS = 256
TAIL_TOKENS = 256
NV_SIZE = 18
N_TARGETS = 8
DEVICE = "cpu"  # keep on CPU for serving

PHQ_ITEM_NAMES = [
    "NoInterest",
    "Depressed",
    "Sleep",
    "Tired",
    "Appetite",
    "Failure",
    "Concentration",
    "Psychomotor",
]

NEGATIVE_WORDS = {
    "hopeless", "worthless", "empty", "tired", "sad", "lonely",
    "fail", "useless", "burden", "numb", "anxious", "hate",
    "depressed", "miserable", "terrible", "awful", "horrible",
    "struggling", "suffering", "pain", "hurt", "crying",
    "helpless", "frustrated", "angry", "guilty", "ashamed",
    "exhausted", "overwhelmed", "stressed", "worried", "afraid",
}

POSITIVE_WORDS = {
    "happy", "good", "fine", "great", "better", "enjoy",
    "love", "hope", "okay", "well", "improving", "grateful",
    "excited", "motivated", "confident", "peaceful", "calm",
    "cheerful", "satisfied", "content", "wonderful", "amazing",
}

FILLER_PATTERN = re.compile(r"\b(um|uh|uhm|hmm|hm|erm|er|ah|like|you know)\b", re.I)


# ═══════════════════════════════════════════════════════════════════════════════
# 2. MODEL DEFINITION  (must match the notebook exactly)
# ═══════════════════════════════════════════════════════════════════════════════

class MPNetMultiOutputRegressor(nn.Module):
    """
    MPNet encoder with attention pooling + 8-head multi-output regressor.

    Architecture:
        Encoder   → last_hidden_state  (B, L, 768)
        AttnPool  → pooled             (B, 768)
        Concat    → [pooled, nv]       (B, 768 + NV_SIZE)
        SharedMLP → shared_repr        (B, 256)
        Heads     → [h1..h8]           8 × Linear(256 → 1)
    """

    def __init__(self, model_name=MODEL_NAME, nv_size=NV_SIZE,
                 n_targets=N_TARGETS, dropout=0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden = self.encoder.config.hidden_size  # 768

        # Attention pooling
        self.token_attn = nn.Linear(hidden, 1)

        # Shared representation
        self.shared_mlp = nn.Sequential(
            nn.Linear(hidden + nv_size, 512),
            nn.LayerNorm(512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(dropout),
        )

        # 8 independent heads for each PHQ-8 item
        self.heads = nn.ModuleList([nn.Linear(256, 1) for _ in range(n_targets)])
        self.dropout = nn.Dropout(dropout)

    def forward(self, input_ids, attention_mask, nv_features):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden = out.last_hidden_state  # (B, L, H)

        # Attention pooling
        scores = self.token_attn(hidden).squeeze(-1)  # (B, L)
        scores = scores.masked_fill(attention_mask == 0, -1e9)
        weights = torch.softmax(scores, dim=1).unsqueeze(-1)  # (B, L, 1)
        pooled = (hidden * weights).sum(dim=1)  # (B, H)
        pooled = self.dropout(pooled)

        # Concat NV features + shared MLP
        x = torch.cat([pooled, nv_features], dim=-1)  # (B, H+NV)
        x = self.shared_mlp(x)  # (B, 256)

        # Multi-output heads
        preds = torch.cat([head(x) for head in self.heads], dim=-1)  # (B, 8)
        return preds


# ═══════════════════════════════════════════════════════════════════════════════
# 3. TOKENIZATION (Head-Tail truncation — matches notebook)
# ═══════════════════════════════════════════════════════════════════════════════

def tokenize_head_tail(text, tokenizer, max_len=MAX_LEN,
                       head=HEAD_TOKENS, tail=TAIL_TOKENS):
    """Tokenize with head-tail truncation strategy."""
    ids = tokenizer(
        text,
        add_special_tokens=False,
        return_attention_mask=False,
        return_token_type_ids=False,
    )["input_ids"]

    bos = tokenizer.cls_token_id or tokenizer.bos_token_id
    eos = tokenizer.sep_token_id or tokenizer.eos_token_id
    inner = max_len - 2

    if len(ids) > inner:
        h = min(head, inner // 2)
        t = inner - h
        ids = ids[:h] + ids[-t:]

    ids = [bos] + ids + [eos]
    pad = max_len - len(ids)
    mask = [1] * len(ids) + [0] * pad
    ids = ids + [tokenizer.pad_token_id] * pad

    return ids, mask


# ═══════════════════════════════════════════════════════════════════════════════
# 4. NV FEATURE EXTRACTION (adapted for chatbot text — no audio)
# ═══════════════════════════════════════════════════════════════════════════════

def extract_nv_features(text: str) -> np.ndarray:
    """
    Extract the 18-dim NV feature vector from text input.

    Since we're processing typed text (not DAIC-WOZ audio transcripts),
    most nonverbal features will be zero. The lexical features (neg_ratio,
    pos_ratio, sentiment_gap) and text length features are still informative.

    Feature order (must match training):
        [nv_laugh, nv_sigh, nv_cough, nv_breath, nv_sniff, nv_groan,
         nv_pause, nv_um, nv_uh, nv_total, filler_count,
         word_count, utterance_count, avg_utt_len, response_brevity,
         neg_ratio, pos_ratio, sentiment_gap]
    """
    words = text.lower().split()
    word_set = set(words)
    total_words = len(words) + 1  # avoid div-by-zero

    # Nonverbal tags: all zero for typed text (no <laugh>, <sigh> etc.)
    nv_tags = [0.0] * 9  # laugh, sigh, cough, breath, sniff, groan, pause, um, uh
    nv_total = 0.0

    # Filler words
    filler_count = len(FILLER_PATTERN.findall(text))

    # Length features
    sentences = [s.strip() for s in re.split(r"[.!?\n]+", text) if s.strip()]
    word_count = len(words)
    utterance_count = max(len(sentences), 1)
    avg_utt_len = word_count / utterance_count
    response_brevity = 1.0 if utterance_count < 50 else 0.0

    # Lexical features
    neg_count = len(word_set & NEGATIVE_WORDS)
    pos_count = len(word_set & POSITIVE_WORDS)
    neg_ratio = neg_count / total_words
    pos_ratio = pos_count / total_words
    sentiment_gap = (neg_count - pos_count) / total_words

    features = np.array(
        nv_tags + [nv_total, filler_count,
                   word_count, utterance_count, avg_utt_len, response_brevity,
                   neg_ratio, pos_ratio, sentiment_gap],
        dtype=np.float32,
    )
    return features  # shape (18,)


def severity_category(score: float) -> dict:
    """Return severity category info based on PHQ-8 total score."""
    score = round(score)
    if score <= 4:
        return {"label": "Minimal", "color": "green", "description": "Minimal depression — no treatment typically needed."}
    elif score <= 9:
        return {"label": "Mild", "color": "amber", "description": "Mild depression — watchful waiting; repeat screening at follow-up."}
    elif score <= 14:
        return {"label": "Moderate", "color": "orange", "description": "Moderate depression — consultation with a mental health professional is recommended."}
    elif score <= 19:
        return {"label": "Moderately Severe", "color": "red-orange", "description": "Moderately severe depression — active treatment with pharmacotherapy and/or psychotherapy is recommended."}
    else:
        return {"label": "Severe", "color": "red", "description": "Severe depression — immediate initiation of pharmacotherapy and referral to a mental health specialist."}


# ═══════════════════════════════════════════════════════════════════════════════
# 5. LOAD MODEL & TOKENIZER
# ═══════════════════════════════════════════════════════════════════════════════

print(f"Loading tokenizer: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print(f"Building model architecture...")
model = MPNetMultiOutputRegressor(
    model_name=MODEL_NAME,
    nv_size=NV_SIZE,
    n_targets=N_TARGETS,
    dropout=0.0,  # no dropout at inference
)

print(f"Loading weights from: {MODEL_PATH}")
state_dict = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)
model.load_state_dict(state_dict)
model.eval()
model.to(DEVICE)
print("✅ Model loaded successfully!")


# ═══════════════════════════════════════════════════════════════════════════════
# 6. FLASK APP
# ═══════════════════════════════════════════════════════════════════════════════

app = Flask(__name__)
CORS(app)


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok", "model": MODEL_NAME, "targets": N_TARGETS})


@app.route("/predict", methods=["POST"])
def predict():
    """
    Predict PHQ-8 scores from 8 interview answers.

    Request body:
        {"answers": ["answer1", "answer2", ..., "answer8"]}

    Response:
        {
            "scores": [0-3, ...],       # 8 per-item scores
            "total_score": 0-24,         # sum of items
            "severity": {...},           # category info
            "filler_count": int,
            "negative_words": [str, ...],
            "positive_words": [str, ...],
            "item_names": [str, ...]
        }
    """
    data = request.get_json(force=True)
    answers = data.get("answers", [])

    if len(answers) != 8:
        return jsonify({"error": "Exactly 8 answers required"}), 400

    # Combine all answers into one transcript (DAIC-WOZ style)
    transcript = " ".join(answers)

    # Tokenize
    ids, mask = tokenize_head_tail(transcript, tokenizer)
    input_ids = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    attention_mask = torch.tensor([mask], dtype=torch.long, device=DEVICE)

    # Extract NV features
    nv = extract_nv_features(transcript)
    nv_tensor = torch.tensor([nv], dtype=torch.float, device=DEVICE)

    # Inference
    with torch.no_grad():
        preds = model(input_ids, attention_mask, nv_tensor)  # (1, 8)

    # Post-process
    scores_raw = preds[0].cpu().numpy()
    scores_clipped = np.clip(scores_raw, 0, 3).tolist()
    scores_rounded = [round(s, 2) for s in scores_clipped]
    total_score = round(sum(scores_rounded), 2)
    severity = severity_category(total_score)

    # Detect filler and sentiment words in the transcript
    words_lower = set(transcript.lower().split())
    detected_neg = sorted(words_lower & NEGATIVE_WORDS)
    detected_pos = sorted(words_lower & POSITIVE_WORDS)
    filler_count = len(FILLER_PATTERN.findall(transcript))

    return jsonify({
        "scores": scores_rounded,
        "total_score": total_score,
        "severity": severity,
        "filler_count": filler_count,
        "negative_words": detected_neg,
        "positive_words": detected_pos,
        "item_names": PHQ_ITEM_NAMES,
    })


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000, debug=False)