import os
import json
import streamlit as st
import faiss
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
from reportlab.lib.pagesizes import A4
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
from reportlab.lib.styles import getSampleStyleSheet

# Load milestones data
with open('milestones.json', 'r') as f:
    milestones = json.load(f)

# Age categories for dropdown selection
age_categories = {
    "Up to 2 months": 2,
    "Up to 4 months": 4,
    "Up to 6 months": 6,
    "Up to 9 months": 9,
    "Up to 1 year": 12,
    "Up to 15 months": 15,
    "Up to 18 months": 18,
    "Up to 2 years": 24,
    "Up to 30 months": 30,
    "Up to 3 years": 36,
    "Up to 4 years": 48,
    "Up to 5 years": 60
}

# Initialize FAISS and Sentence Transformer
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_faiss_index(data):
    descriptions = []
    age_keys = []
    for age, categories in data.items():
        for entry in categories:
            descriptions.append(entry['description'])
            age_keys.append(int(age))
    
    embeddings = model.encode(descriptions, convert_to_numpy=True)
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index, descriptions, age_keys

index, descriptions, age_keys = create_faiss_index(milestones)

# Function to retrieve the closest milestone
def retrieve_milestone(user_input):
    user_embedding = model.encode([user_input], convert_to_numpy=True)
    _, indices = index.search(user_embedding, 1)
    return descriptions[indices[0][0]] if indices[0][0] < len(descriptions) else "No relevant milestone found."

# Initialize IBM Granite Model
BASE_NAME = "ibm-granite/granite-3.0-2b-base"

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

tokenizer = AutoTokenizer.from_pretrained(BASE_NAME, padding_side='left', trust_remote_code=True)
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME, device_map="auto")

def generate_response(user_input, child_age):
    relevant_milestone = retrieve_milestone(user_input)
    question_chat = [
        {
            "role": "system",
            "content": f"The child is {child_age} months old. Based on the given traits: {user_input}, determine whether the child is meeting expected milestones. Relevant milestone: {relevant_milestone}. If there are any concerns, suggest steps the parents can take."
        },
        {
            "role": "user",
            "content": user_input
        }
    ]
    input_text = tokenizer.apply_chat_template(question_chat, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt")
    output = model_base.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=500)
    output_text = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return output_text

# Streamlit UI Styling
st.set_page_config(page_title="Tiny Triumphs Tracker", page_icon="👶", layout="wide")
st.markdown("""
    <style>
        .stApp { background-color: #1e1e2e; color: #ffffff; }
        .stTitle { text-align: center; color: #ffcc00; font-size: 36px; font-weight: bold; }
    </style>
""", unsafe_allow_html=True)

st.markdown("<h1 class='stTitle'>👶 Tiny Triumphs Tracker</h1>", unsafe_allow_html=True)
st.markdown("Track your child's key growth milestones from birth to 5 years and detect early developmental concerns.", unsafe_allow_html=True)

# User selects child's age
selected_age = st.selectbox("📅 Select child's age:", list(age_categories.keys()))
child_age = age_categories[selected_age]

# User input for traits and skills
placeholder_text = "Describe your child's behavior and skills."
user_input = st.text_area("✍️ Enter child's behavioral traits and skills:", placeholder=placeholder_text)

def generate_pdf_report(ai_response):
    pdf_file = "progress_report.pdf"
    doc = SimpleDocTemplate(pdf_file, pagesize=A4)
    styles = getSampleStyleSheet()
    elements = []
    elements.append(Paragraph("Child Development Progress Report", styles['Title']))
    elements.append(Spacer(1, 12))
    elements.append(Paragraph("Development Insights:", styles['Heading2']))
    elements.append(Spacer(1, 10))
    response_parts = ai_response.split('\n')
    for part in response_parts:
        part = part.strip().lstrip('0123456789.- ')
        if part:
            elements.append(Paragraph(f"• {part}", styles['Normal']))
            elements.append(Spacer(1, 5))
    disclaimer = "This report is AI-generated and is for informational purposes only. "
    elements.append(Spacer(1, 12))
    elements.append(Paragraph(disclaimer, styles['Italic']))
    doc.build(elements)
    return pdf_file

if st.button("🔍 Analyze", help="Click to analyze the child's development milestones"):
    ai_response = generate_response(user_input, child_age)
    st.subheader("📊 Development Insights:")
    st.markdown(f"<div style='background-color:#44475a; color:#ffffff; padding: 15px; border-radius: 10px;'>{ai_response}</div>", unsafe_allow_html=True)
    pdf_file = generate_pdf_report(ai_response)
    with open(pdf_file, "rb") as f:
        st.download_button(label="📥 Download Progress Report", data=f, file_name="progress_report.pdf", mime="application/pdf")

st.warning("⚠️ The results provided are generated by AI and should be interpreted with caution. Please consult a pediatrician for professional advice.")