"""
============================================================
Mentallico AI v2.0 - Central Configuration
============================================================
مركز الإعدادات لكل مكونات النظام
"""
import os
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

# ============= Paths =============
BASE_DIR = Path(__file__).parent.resolve()
DATA_DIR = BASE_DIR / "data"
MODELS_DIR = BASE_DIR / "saved_models"
CHROMA_DB_PATH = str(BASE_DIR / "chroma_db")
PDF_DOCS_DIR = str(BASE_DIR / "medical_docs")

MODELS_DIR.mkdir(exist_ok=True)
DATA_DIR.mkdir(exist_ok=True)
Path(PDF_DOCS_DIR).mkdir(exist_ok=True)

# ============= Classification Models =============
# أصبح لدينا 4 نماذج تعمل معاً (Ensemble)
CLASSIFICATION_MODELS = {
    "primary": {
        "name": "mental/mental-bert-base-uncased",
        "fallback": "distilbert-base-uncased",
        "weight": 0.35, # تم تعديل الأوزان لتقسيمها على 4 نماذج
        "description": "BERT مدرب خصيصاً على نصوص الصحة النفسية"
    },
    "secondary": {
        "name": "distilroberta-base",
        "fallback": "distilbert-base-uncased",
        "weight": 0.25,
        "description": "RoBERTa - قوي في فهم السياق"
    },
    "tertiary": {
        "name": "dsuram/distilbert-mentalhealth-classifier",
        "fallback": None, # Zero-shot/Pipeline
        "weight": 0.25,
        "description": "نموذج جاهز للصحة النفسية - يدعم التنبؤات الأساسية"
    },
    "emotion": {
        "name": "j-hartmann/emotion-english-distilroberta-base",
        "fallback": None,
        "weight": 0.15,
        "description": "تحليل المشاعر"
    }
}

# الفئات المرضية الأساسية
LABELS = ['Normal', 'Depression', 'Suicide', 'Anxiety/Stress', 'Bipolar', 'PTSD']

# ============= Training Config =============
TRAIN_CONFIG = {
    'max_len': 128,
    'batch_size': 32,
    'epochs': 5,
    'lr': 2e-5,
    'weight_decay': 0.01,
    'warmup_ratio': 0.1,
    'use_focal_loss': True,
    'focal_alpha': 0.25,
    'focal_gamma': 2.0,
    'bipolar_weight': 2.0,
    'bipolar_target': 1500,
    'val_size': 0.2,
    'test_size': 0.2,
    'random_seed': 42,
    'early_stopping_patience': 3,
    'dropout': 0.3,
}

# ============= Diagnosis Logic =============
DIAGNOSIS_CONFIG = {
    'min_messages_before_diagnosis': 3,
    'min_total_words': 60,
    'confidence_threshold': 0.30,
    'consistency_window': 5,
    'consistency_threshold': 0.66,
    'high_risk_immediate': ['Suicide'],
    'max_history_messages': 50,
}

# ============= RAG Config =============
RAG_CONFIG = {
    'embedding_model': 'all-MiniLM-L6-v2',
    'llm_model': 'llama-3.3-70b-versatile',  # تم التغيير إلى Gemini (يفضل pro للتشخيص الطبي)
    'temperature': 0.3,
    'chunk_size': 2500,
    'chunk_overlap': 500,
    'retrieval_k': 12,
    'max_tokens': 1024,
}

# ============= STT Config =============
STT_CONFIG = {
    'model': 'whisper-large-v3',
    'language': None,
    'fallback_to_local': True,
    'local_model_size': 'base',
}

# ============= API Keys =============
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "") # أبقيناه في حال احتجته للـ STT

# ============= Severity Levels =============
SEVERITY_LEVELS = {
    'Normal': 0,
    'Anxiety/Stress': 1,
    'Depression': 2,
    'PTSD': 2,
    'Bipolar': 3,
    'Suicide': 4,
}