from fastapi import FastAPI
from pydantic import BaseModel
import joblib
from transformers import pipeline
import re

app = FastAPI(title="🩺 MediAgent AI API (No CSV Version)")

# =========================
# LOAD MODEL
# =========================
data = joblib.load("diseases.pkl")

model = data["model"]
encoder = data["encoder"]
symptoms_list = data["symptoms"]

# =========================
# NORMALIZATION
# =========================
def normalize(text):
    text = str(text).lower()
    text = re.sub(r"\(.*?\)", "", text)
    text = re.sub(r"[^a-z0-9\s]", " ", text)
    return " ".join(text.split())

# =========================
# LANGUAGE (HINGLISH SUPPORT)
# =========================
def hinglish_to_english(text):
    mapping = {
        "bukhar": "fever",
        "khansi": "cough",
        "sar dard": "headache",
        "saans": "breathing",
        "dard": "pain"
    }
    for h, e in mapping.items():
        text = text.lower().replace(h, e)
    return text.lower()

# =========================
# NLP MODEL (LAZY LOAD)
# =========================
classifier = None

def get_classifier():
    global classifier
    if classifier is None:
        classifier = pipeline(
            "zero-shot-classification",
            model="typeform/distilbert-base-uncased-mnli"
        )
    return classifier

labels = ["fever", "cough", "headache", "vomiting", "chest pain", "breathlessness"]

label_map = {
    "fever": "fever",
    "cough": "cough",
    "headache": "headache",
    "vomiting": "vomiting",
    "chest pain": "chest_pain",
    "breathlessness": "breathlessness"
}

# =========================
# SYMPTOM EXTRACTION
# =========================
def extract_symptoms(text):
    clf = get_classifier()
    result = clf(text, labels)

    symptoms = []
    for label, score in zip(result["labels"], result["scores"]):
        if score > 0.4:
            symptoms.append(label_map[label])

    text = text.lower()

    if "breath" in text:
        symptoms.append("breathlessness")
    if "chest" in text and "pain" in text:
        symptoms.append("chest_pain")
    if "fever" in text:
        symptoms.append("fever")
    if "cough" in text:
        symptoms.append("cough")
    if "vomit" in text:
        symptoms.append("vomiting")
    if "head" in text and "pain" in text:
        symptoms.append("headache")

    return list(set(symptoms))

# =========================
# SEVERITY (NO CSV VERSION)
# =========================
def get_severity(symptoms):
    score = len(symptoms)

    if score >= 4:
        return "Critical"
    elif score == 3:
        return "High"
    elif score == 2:
        return "Moderate"
    return "Low"

# =========================
# DISEASE PREDICTION
# =========================
def predict_disease(symptoms):
    vec = [1 if s in symptoms else 0 for s in symptoms_list]
    pred = model.predict([vec])[0]
    return encoder.inverse_transform([pred])[0]

# =========================
# DESCRIPTION GENERATOR
# =========================
def generate_description(disease):
    return f"{disease} is a medical condition that requires proper diagnosis and treatment. Consult a healthcare professional for detailed evaluation."

# =========================
# PRECAUTION GENERATOR (FULL COVERAGE)
# =========================
def generate_precautions(disease):
    d = normalize(disease)

    # INFECTIOUS
    if any(x in d for x in ["infection", "fever", "virus", "bacteria", "abscess", "tuberculosis", "dengue", "malaria"]):
        return [
            "Maintain personal hygiene",
            "Wash hands frequently",
            "Avoid contact with infected individuals",
            "Use clean water and food",
            "Follow vaccination guidelines"
        ]

    # RESPIRATORY
    elif any(x in d for x in ["lung", "asthma", "bronch", "respiratory", "pneumonia"]):
        return [
            "Avoid smoking and pollution",
            "Use masks in dusty environments",
            "Avoid allergens",
            "Practice breathing exercises",
            "Take prescribed medication"
        ]

    # HEART
    elif any(x in d for x in ["heart", "cardio", "hypertension", "artery"]):
        return [
            "Eat low-salt, low-fat diet",
            "Exercise regularly",
            "Avoid smoking",
            "Manage stress",
            "Monitor blood pressure"
        ]

    # METABOLIC
    elif any(x in d for x in ["diabetes", "thyroid", "obesity"]):
        return [
            "Maintain balanced diet",
            "Exercise regularly",
            "Monitor sugar levels",
            "Avoid processed food",
            "Follow medication plan"
        ]

    # MENTAL
    elif any(x in d for x in ["depression", "anxiety", "mental", "psych"]):
        return [
            "Maintain sleep schedule",
            "Practice relaxation techniques",
            "Seek counseling",
            "Avoid alcohol",
            "Stay socially active"
        ]

    # SKIN
    elif any(x in d for x in ["skin", "acne", "eczema", "dermatitis", "fungal"]):
        return [
            "Maintain hygiene",
            "Avoid irritants",
            "Keep skin dry",
            "Use prescribed creams",
            "Avoid sharing personal items"
        ]

    # INJURY / BONES
    elif any(x in d for x in ["fracture", "injury", "arthritis", "joint", "bone"]):
        return [
            "Avoid heavy strain",
            "Use protective gear",
            "Ensure calcium intake",
            "Follow physiotherapy",
            "Take rest"
        ]

    # CANCER
    elif "cancer" in d:
        return [
            "Avoid smoking",
            "Eat healthy diet",
            "Regular screening",
            "Avoid harmful chemicals",
            "Consult doctor early"
        ]

    # DEFAULT
    else:
        return [
            "Maintain healthy lifestyle",
            "Eat balanced diet",
            "Exercise regularly",
            "Avoid smoking and alcohol",
            "Consult doctor if symptoms worsen"
        ]

# =========================
# REQUEST MODEL
# =========================
class InputData(BaseModel):
    text: str

# =========================
# ROUTES
# =========================
@app.get("/")
def home():
    return {"message": "MediAgent API running (No CSV) 🚀"}

@app.post("/predict")
def predict_api(input: InputData):
    text = input.text

    if not text:
        return {"error": "No input provided"}

    original = text
    text = hinglish_to_english(text)

    symptoms = extract_symptoms(text)

    if "सांस" in original:
        symptoms.append("breathlessness")

    symptoms = list(set(symptoms))

    if not symptoms:
        return {"error": "No symptoms detected"}

    disease = predict_disease(symptoms)

    description = generate_description(disease)
    precautions = generate_precautions(disease)
    severity = get_severity(symptoms)

    return {
        "symptoms": symptoms,
        "disease": disease,
        "severity": severity,
        "description": description,
        "precautions": precautions
    }