apps/streamlit_ui/pages/1_Profile.py

"""User profile management page."""

import sys
from pathlib import Path

# Add the project root to the Python path
# This is necessary for Streamlit to find modules in the 'apps' directory
project_root = Path(__file__).resolve().parents[3]
if str(project_root) not in sys.path:
    sys.path.append(str(project_root))

import pandas as pd
import streamlit as st

from apps.streamlit_ui.page_versions.profile import v1, v2
from sentinel.models import (
    ClinicalObservation,
    Demographics,
    FamilyMemberCancer,
    FemaleSpecific,
    Lifestyle,
    PersonalMedicalHistory,
    UserInput,
)
from sentinel.utils import load_user_file


# --- Helper Functions ---
def clear_profile_state():
    """Callback function to reset profile-related session state."""
    st.session_state.user_profile = None
    if "profile_upload" in st.session_state:
        del st.session_state["profile_upload"]


# --- Main Page Layout ---
st.title("👤 User Profile")

# --- Sidebar for Version Selection and Upload ---
with st.sidebar:
    st.header("Controls")

    # Version selection
    version_options = ["V2 (Editable Form)", "V1 (JSON Viewer)"]
    version = st.radio(
        "Select Demo Version",
        version_options,
        help="Choose the version of the profile page to display.",
    )

    st.divider()

    # Example Profile Selector
    examples_dir = project_root / "examples"

    # Collect all example profiles
    profile_files = []
    if examples_dir.exists():
        # Get profiles from dev/
        dev_dir = examples_dir / "dev"
        if dev_dir.exists():
            profile_files.extend(sorted(dev_dir.glob("*.yaml")))
            profile_files.extend(sorted(dev_dir.glob("*.json")))

        # Get profiles from synthetic/
        synthetic_dir = examples_dir / "synthetic"
        if synthetic_dir.exists():
            for subdir in sorted(synthetic_dir.iterdir()):
                if subdir.is_dir():
                    profile_files.extend(sorted(subdir.glob("*.yaml")))
                    profile_files.extend(sorted(subdir.glob("*.json")))

    # Create display names (relative to examples/)
    profile_options = {}
    if profile_files:
        for p in profile_files:
            rel_path = p.relative_to(examples_dir)
            profile_options[str(rel_path)] = p

    # Dropdown selector
    if profile_options:
        selected = st.selectbox(
            "Load Example Profile",
            options=["-- Select a profile --", *profile_options.keys()],
            key="profile_selector",
        )

        if selected != "-- Select a profile --":
            try:
                profile_path = profile_options[selected]
                st.session_state.user_profile = load_user_file(str(profile_path))
                st.success(f"✅ Loaded: {selected}")
            except Exception as e:
                st.error(f"Failed to load profile: {e}")

    # Clear Profile Button
    if st.session_state.get("user_profile"):
        st.button(
            "Clear Loaded Profile",
            on_click=clear_profile_state,
            use_container_width=True,
        )


# --- Page Content Dispatcher ---
# Render the selected page version
if version == "V1 (JSON Viewer)":
    v1.render()
else:  # Default to V2
    v2.render()

# The manual creation form can be a persistent feature at the bottom of the page
with st.expander("Create New Profile Manually"):
    # --- STEP 1: Move the sex selector OUTSIDE the form. ---
    # This allows it to trigger a rerun and update the UI dynamically.
    # Give it a unique key to avoid conflicts with other widgets.
    sex = st.selectbox(
        "Biological Sex", ["Male", "Female", "Other"], key="manual_profile_sex"
    )

    with st.form("manual_profile_form"):
        st.subheader("Demographics")
        age = st.number_input("Age", min_value=0, step=1)
        # The 'sex' variable is now taken from the selector above the form.
        ethnicity = st.text_input("Ethnicity")

        st.subheader("Lifestyle")
        smoking_status = st.selectbox("Smoking Status", ["never", "former", "current"])
        smoking_pack_years = st.number_input("Pack-Years", min_value=0, step=1)
        alcohol_consumption = st.selectbox(
            "Alcohol Consumption", ["none", "light", "moderate", "heavy"]
        )
        dietary_habits = st.text_area("Dietary Habits")
        physical_activity_level = st.text_area("Physical Activity")

        st.subheader("Personal Medical History")
        known_genetic_mutations = st.text_input(
            "Known Genetic Mutations (comma-separated)"
        )
        previous_cancers = st.text_input("Previous Cancers (comma-separated)")
        chronic_illnesses = st.text_input("Chronic Illnesses (comma-separated)")

        st.subheader("Family History")
        fam_cols = ["relative", "cancer_type", "age_at_diagnosis"]
        fam_df = st.data_editor(
            pd.DataFrame(columns=fam_cols),
            num_rows="dynamic",
            key="family_history_editor",
        )

        st.subheader("Clinical Observations")
        obs_cols = ["test_name", "value", "unit", "reference_range", "date"]
        obs_df = st.data_editor(
            pd.DataFrame(columns=obs_cols),
            num_rows="dynamic",
            key="clinical_obs_editor",
        )

        female_specific_data = {}
        # --- STEP 2: The conditional check now works correctly. ---
        # The 'if' statement is evaluated on each rerun when the 'sex' selector changes.
        if sex == "Female":
            st.subheader("Female-Specific")
            female_specific_data["age_at_first_period"] = st.number_input(
                "Age at First Period", min_value=0, step=1
            )
            female_specific_data["age_at_menopause"] = st.number_input(
                "Age at Menopause", min_value=0, step=1
            )
            female_specific_data["num_live_births"] = st.number_input(
                "Number of Live Births", min_value=0, step=1
            )
            female_specific_data["age_at_first_live_birth"] = st.number_input(
                "Age at First Live Birth", min_value=0, step=1
            )
            female_specific_data["hormone_therapy_use"] = st.text_input(
                "Hormone Therapy Use"
            )

        current_concerns = st.text_area("Current Concerns or Symptoms")

        submitted = st.form_submit_button("Save New Profile")
        if submitted:
            # --- STEP 3: Use the 'sex' variable from the external selector during submission. ---
            demographics = Demographics(
                age=int(age), sex=sex, ethnicity=ethnicity or None
            )
            lifestyle = Lifestyle(
                smoking_status=smoking_status,
                smoking_pack_years=int(smoking_pack_years) or None,
                alcohol_consumption=alcohol_consumption,
                dietary_habits=dietary_habits or None,
                physical_activity_level=physical_activity_level or None,
            )
            pmh = PersonalMedicalHistory(
                known_genetic_mutations=[
                    m.strip() for m in known_genetic_mutations.split(",") if m.strip()
                ],
                previous_cancers=[
                    c.strip() for c in previous_cancers.split(",") if c.strip()
                ],
                chronic_illnesses=[
                    i.strip() for i in chronic_illnesses.split(",") if i.strip()
                ],
            )
            family_history = []
            for _, row in fam_df.dropna(how="all").iterrows():
                if row.get("relative") and row.get("cancer_type"):
                    family_history.append(
                        FamilyMemberCancer(
                            relative=str(row["relative"]),
                            cancer_type=str(row["cancer_type"]),
                            age_at_diagnosis=int(row["age_at_diagnosis"])
                            if row["age_at_diagnosis"] not in ["", None]
                            else None,
                        )
                    )

            observations = []
            for _, row in obs_df.dropna(how="all").iterrows():
                if row.get("test_name") and row.get("value") and row.get("unit"):
                    observations.append(
                        ClinicalObservation(
                            test_name=str(row["test_name"]),
                            value=str(row["value"]),
                            unit=str(row["unit"]),
                            reference_range=(
                                str(row["reference_range"])
                                if row["reference_range"] not in ["", None]
                                else None
                            ),
                            date=str(row["date"])
                            if row["date"] not in ["", None]
                            else None,
                        )
                    )

            female_specific = None
            if sex == "Female":
                female_specific = FemaleSpecific(**female_specific_data)

            new_profile = UserInput(
                demographics=demographics,
                lifestyle=lifestyle,
                family_history=family_history,
                personal_medical_history=pmh,
                female_specific=female_specific,
                current_concerns_or_symptoms=current_concerns or None,
                clinical_observations=observations,
            )
            st.success("Profile saved")

            # --- STEP 4: Compute the risk scores ---
            with st.spinner("Calculating risk scores..."):
                from sentinel.risk_models import RISK_MODELS

                risks_scores = []
                for model in RISK_MODELS:
                    risk_score = model().run(new_profile)
                    # Handle models that return multiple scores (e.g., QCancer)
                    if isinstance(risk_score, list):
                        risks_scores.extend(risk_score)
                    else:
                        risks_scores.append(risk_score)

                new_profile.risks_scores = risks_scores

            st.session_state.user_profile = new_profile
            st.success("Risk scores calculated!")
            st.rerun()