"""
generate_seed.py — Realistic seed data for Kasualdad LFED.

Generates 5 schools × 4 school years × 12 grade levels with:
  - enrollment: aggregate student counts per school/year/grade
  - attendance: student-level absence records
  - students: demographics (gender, race_ethnicity, ELL, SpEd, FRL)
  - discipline: incident records (type, severity, action, days suspended)
  - grades: course grades + GPA per term

Attendance is the focal domain. Other tables provide context for
cross-table queries (e.g., "chronically absent ELL students",
"discipline incidents for students with low GPA").

Run standalone:  python data/generate_seed.py
Used by:         data_engine.seed_database()
"""

import random
import duckdb
from pathlib import Path

# ── Reproducibility ────────────────────────────────────────────────────

SEED = 42
random.seed(SEED)

# ── School definitions ─────────────────────────────────────────────────

SCHOOLS = [
    {
        "name": "Lincoln Elementary",
        "grades": list(range(0, 6)),   # K-5
        "base_enrollment": 520,
        "students": [],
    },
    {
        "name": "Washington Middle",
        "grades": list(range(6, 9)),   # 6-8
        "base_enrollment": 480,
        "students": [],
    },
    {
        "name": "Jefferson High",
        "grades": list(range(9, 13)),  # 9-12
        "base_enrollment": 900,
        "students": [],
    },
    {
        "name": "Roosevelt Academy",
        "grades": list(range(0, 9)),   # K-8
        "base_enrollment": 380,
        "students": [],
    },
    {
        "name": "Kennedy Prep",
        "grades": list(range(6, 13)),  # 6-12
        "base_enrollment": 620,
        "students": [],
    },
]

SCHOOL_YEARS = ["2021-2022", "2022-2023", "2023-2024", "2024-2025"]

# ── Distribution parameters ────────────────────────────────────────────

CHRONIC_ABSENT_RATE = 0.15   # ~15% of students
YOY_GROWTH = 0.03            # 3% enrollment growth per year
GRADE_SIZE_VARIANCE = 0.15   # ±15% random variance per grade

# Absence distributions
ABSENCES_NORMAL_MEAN = 5.0
ABSENCES_NORMAL_STD = 3.0
ABSENCES_CHRONIC_MIN = 18
ABSENCES_CHRONIC_MAX = 45

# Demographics (roughly CA-aligned)
GENDERS = ["F", "M"]
GENDER_WEIGHTS = [0.49, 0.51]

RACE_ETHNICITIES = [
    "Hispanic/Latino", "White", "Black/African American",
    "Asian", "Filipino", "Two or More Races", "American Indian",
    "Pacific Islander",
]
RACE_WEIGHTS = [
    0.55, 0.22, 0.06, 0.06, 0.04, 0.05, 0.01, 0.01,
]

ELL_RATE = 0.18
SPED_RATE = 0.12
FRL_RATE = 0.45  # economically disadvantaged

# Discipline
INCIDENT_TYPES = [
    "Defiance", "Fighting", "Vandalism", "Bullying",
    "Disruption", "Insubordination", "Theft", "Substance",
    "Harassment", "Verbal Altercation",
]
INCIDENT_TYPE_WEIGHTS = [
    0.20, 0.15, 0.08, 0.10,
    0.12, 0.10, 0.05, 0.05,
    0.08, 0.07,
]
SEVERITIES = ["Minor", "Major", "Severe"]
SEVERITY_WEIGHTS = [0.70, 0.25, 0.05]
ACTIONS = ["Warning", "Detention", "Suspension", "Expulsion"]
ACTION_WEIGHTS = [0.55, 0.25, 0.18, 0.02]

# Grades
COURSE_NAMES = [
    "English", "Math", "Science", "Social Studies",
    "PE", "Art", "Music",
]
LETTER_GRADES = ["A", "B", "C", "D", "F"]
GRADE_DISTRIBUTION = [0.20, 0.35, 0.30, 0.10, 0.05]  # A, B, C, D, F
TERMS = ["Fall", "Spring"]


# ── Student generation ─────────────────────────────────────────────────

def _generate_students() -> list[dict]:
    """
    Generate ~2,900 students distributed across schools proportional to enrollment.
    Each student gets: id, school_name, grade_level, is_chronically_absent,
                       gender, race_ethnicity, ell/sped/frl flags.
    """
    students = []
    student_id = 1000

    for school in SCHOOLS:
        num_students = school["base_enrollment"]
        for _ in range(num_students):
            grade = random.choice(school["grades"])
            student = {
                "student_id": student_id,
                "school_name": school["name"],
                "grade_level": grade,
                "is_chronically_absent": False,  # set below
                "gender": random.choices(GENDERS, weights=GENDER_WEIGHTS)[0],
                "race_ethnicity": random.choices(
                    RACE_ETHNICITIES, weights=RACE_WEIGHTS
                )[0],
                "english_learner": random.random() < ELL_RATE,
                "special_education": random.random() < SPED_RATE,
                "economically_disadvantaged": random.random() < FRL_RATE,
            }
            students.append(student)
            student_id += 1

    # Mark ~15% as chronically absent
    num_chronic = int(len(students) * CHRONIC_ABSENT_RATE)
    chronic_indices = random.sample(range(len(students)), num_chronic)
    for idx in chronic_indices:
        students[idx]["is_chronically_absent"] = True

    return students


# ── Enrollment generation ──────────────────────────────────────────────

def _generate_enrollment(students: list[dict]) -> list[tuple]:
    """Generate enrollment rows: one per (school_year, school_name, grade_level)."""
    rows = []

    for school in SCHOOLS:
        name = school["name"]
        for year_idx, year in enumerate(SCHOOL_YEARS):
            growth = (1 + YOY_GROWTH) ** year_idx
            for grade in school["grades"]:
                base_count = sum(
                    1 for s in students
                    if s["school_name"] == name and s["grade_level"] == grade
                )
                variance = 1.0 + random.uniform(-GRADE_SIZE_VARIANCE, GRADE_SIZE_VARIANCE)
                count = max(5, int(base_count * growth * variance))
                rows.append((year, name, grade, count))

    return rows


# ── Attendance generation ──────────────────────────────────────────────

def _generate_attendance(students: list[dict]) -> list[tuple]:
    """Generate attendance rows: one per student per school year."""
    rows = []

    for student in students:
        is_chronic = student["is_chronically_absent"]

        for year in SCHOOL_YEARS:
            if is_chronic:
                absences = random.randint(ABSENCES_CHRONIC_MIN, ABSENCES_CHRONIC_MAX)
            else:
                absences = max(0, min(17, int(random.gauss(ABSENCES_NORMAL_MEAN, ABSENCES_NORMAL_STD))))

            rows.append((
                student["student_id"],
                student["school_name"],
                year,
                absences,
                is_chronic,
            ))

    return rows


# ── Students (demographics) generation ─────────────────────────────────

def _generate_students_table(students: list[dict]) -> list[tuple]:
    """Extract student demographics into a flat table."""
    return [
        (
            s["student_id"],
            s["school_name"],
            s["grade_level"],
            s["gender"],
            s["race_ethnicity"],
            s["english_learner"],
            s["special_education"],
            s["economically_disadvantaged"],
        )
        for s in students
    ]


# ── Discipline generation ──────────────────────────────────────────────

def _generate_discipline(students: list[dict]) -> list[tuple]:
    """
    Generate discipline incidents. ~8% of students have 1+ incident per year.
    Chronically absent students are 2x more likely to have incidents.
    """
    rows = []
    incident_id = 50000

    for student in students:
        is_chronic = student["is_chronically_absent"]
        base_incident_rate = 0.08
        if is_chronic:
            base_incident_rate *= 2.0

        for year_idx, year in enumerate(SCHOOL_YEARS):
            if random.random() < base_incident_rate:
                # 1-3 incidents
                num_incidents = random.choices([1, 2, 3], weights=[0.70, 0.25, 0.05])[0]
                for _ in range(num_incidents):
                    severity = random.choices(SEVERITIES, weights=SEVERITY_WEIGHTS)[0]
                    action = random.choices(ACTIONS, weights=ACTION_WEIGHTS)[0]
                    # Days suspended: 0 for warning/detention, 1-5 for suspension
                    if action == "Suspension":
                        days = random.randint(1, 5)
                    elif action == "Expulsion":
                        days = random.randint(30, 90)
                    else:
                        days = 0

                    # Synthesize an incident date within the school year
                    # School year starts ~Aug 15, ends ~Jun 10
                    month = random.randint(8, 12) if random.random() < 0.6 else random.randint(1, 6)
                    day = random.randint(1, 28)
                    if month >= 8:
                        year_for_date = year.split("-")[0]
                    else:
                        year_for_date = year.split("-")[1]
                    incident_date = f"{year_for_date}-{month:02d}-{day:02d}"

                    rows.append((
                        incident_id,
                        student["student_id"],
                        student["school_name"],
                        year,
                        student["grade_level"],
                        random.choices(INCIDENT_TYPES, weights=INCIDENT_TYPE_WEIGHTS)[0],
                        incident_date,
                        severity,
                        action,
                        days,
                    ))
                    incident_id += 1

    return rows


# ── Grades generation ──────────────────────────────────────────────────

def _generate_grades(students: list[dict]) -> list[tuple]:
    """
    Generate course grades: 2 terms × 7 courses per student per year.
    Chronically absent students skew toward lower grades.
    """
    rows = []
    grade_numeric_map = {"A": 4.0, "B": 3.0, "C": 2.0, "D": 1.0, "F": 0.0}

    for student in students:
        is_chronic = student["is_chronically_absent"]

        # Adjust grade distribution for chronically absent students
        if is_chronic:
            grade_dist = [0.05, 0.15, 0.30, 0.30, 0.20]  # shift lower
        else:
            grade_dist = GRADE_DISTRIBUTION

        for year in SCHOOL_YEARS:
            for term in TERMS:
                term_grades = []
                for course in COURSE_NAMES:
                    letter = random.choices(LETTER_GRADES, weights=grade_dist)[0]
                    numeric = grade_numeric_map[letter]
                    term_grades.append(numeric)
                    rows.append((
                        student["student_id"],
                        student["school_name"],
                        year,
                        student["grade_level"],
                        course,
                        term,
                        letter,
                        numeric,
                        None,  # GPA filled below per term
                    ))

                # Compute term GPA = mean of course numeric grades
                term_gpa = round(sum(term_grades) / len(term_grades), 2)
                # Backfill GPA for this term's rows
                for i in range(-len(COURSE_NAMES), 0):
                    rows[i] = rows[i][:8] + (term_gpa,)

    return rows


# ── Public API ─────────────────────────────────────────────────────────

def generate_seed_data(conn: duckdb.DuckDBPyConnection) -> None:
    """
    Generate and insert all seed data into a DuckDB connection.

    Creates five tables:
      - enrollment(school_year, school_name, grade_level, student_count)
      - attendance(student_id, school_name, school_year, absence_count,
                   is_chronically_absent)
      - students(student_id, school_name, grade_level, gender, race_ethnicity,
                 english_learner, special_education, economically_disadvantaged)
      - discipline(incident_id, student_id, school_name, school_year, grade_level,
                   incident_type, incident_date, severity, action_taken, days_suspended)
      - grades(student_id, school_name, school_year, grade_level, course_name,
               term, letter_grade, grade_numeric, gpa)
    """
    # Re-seed on every call: module-level random.seed(SEED) only fixes the
    # RNG once per process, so repeated calls (e.g. per-request fallback
    # seeding) would otherwise produce different data each time.
    random.seed(SEED)

    # 1. Create tables
    conn.execute("""
        CREATE TABLE IF NOT EXISTS enrollment (
            school_year VARCHAR,
            school_name VARCHAR,
            grade_level INTEGER,
            student_count INTEGER
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS attendance (
            student_id INTEGER,
            school_name VARCHAR,
            school_year VARCHAR,
            absence_count INTEGER,
            is_chronically_absent BOOLEAN
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS students (
            student_id INTEGER,
            school_name VARCHAR,
            grade_level INTEGER,
            gender VARCHAR,
            race_ethnicity VARCHAR,
            english_learner BOOLEAN,
            special_education BOOLEAN,
            economically_disadvantaged BOOLEAN
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS discipline (
            incident_id INTEGER,
            student_id INTEGER,
            school_name VARCHAR,
            school_year VARCHAR,
            grade_level INTEGER,
            incident_type VARCHAR,
            incident_date VARCHAR,
            severity VARCHAR,
            action_taken VARCHAR,
            days_suspended INTEGER
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS grades (
            student_id INTEGER,
            school_name VARCHAR,
            school_year VARCHAR,
            grade_level INTEGER,
            course_name VARCHAR,
            term VARCHAR,
            letter_grade VARCHAR,
            grade_numeric DOUBLE,
            gpa DOUBLE
        )
    """)

    # 2. Generate data
    students = _generate_students()
    enrollment = _generate_enrollment(students)
    attendance = _generate_attendance(students)
    students_table = _generate_students_table(students)
    discipline = _generate_discipline(students)
    grades = _generate_grades(students)

    # 3. Insert
    conn.executemany("INSERT INTO enrollment VALUES (?, ?, ?, ?)", enrollment)
    conn.executemany("INSERT INTO attendance VALUES (?, ?, ?, ?, ?)", attendance)
    conn.executemany("INSERT INTO students VALUES (?, ?, ?, ?, ?, ?, ?, ?)", students_table)
    conn.executemany("INSERT INTO discipline VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", discipline)
    conn.executemany("INSERT INTO grades VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", grades)

    # 4. Stats
    total_students = len(students)
    chronic_count = sum(1 for s in students if s["is_chronically_absent"])
    print(
        f"📊 Seed data: {total_students} students, "
        f"{chronic_count} chronic ({100*chronic_count/total_students:.1f}%), "
        f"{len(enrollment)} enrollment rows, "
        f"{len(attendance)} attendance rows, "
        f"{len(discipline)} discipline rows, "
        f"{len(grades)} grade rows"
    )


# ── Standalone run ─────────────────────────────────────────────────────

if __name__ == "__main__":
    conn = duckdb.connect(":memory:")
    generate_seed_data(conn)

    # Quick verification queries
    print("\n── Sample enrollment ──")
    print(conn.execute(
        "SELECT * FROM enrollment WHERE school_year = '2024-2025' LIMIT 10"
    ).fetchdf())

    print("\n── Chronic absence by school (2023-2024) ──")
    print(conn.execute(
        """
        SELECT
            school_name,
            COUNT(*) AS total_students,
            SUM(CASE WHEN is_chronically_absent THEN 1 ELSE 0 END) AS chronic_count,
            ROUND(100.0 * SUM(CASE WHEN is_chronically_absent THEN 1 ELSE 0 END) / COUNT(*), 1) AS chronic_pct
        FROM attendance
        WHERE school_year = '2023-2024'
        GROUP BY school_name
        ORDER BY school_name
        """
    ).fetchdf())

    print("\n── Chronically absent ELL students (2023-2024) ──")
    print(conn.execute(
        """
        SELECT
            a.school_name,
            COUNT(DISTINCT a.student_id) AS chronic_ell_students
        FROM attendance a
        JOIN students s ON a.student_id = s.student_id
        WHERE a.school_year = '2023-2024'
          AND a.is_chronically_absent = TRUE
          AND s.english_learner = TRUE
        GROUP BY a.school_name
        ORDER BY a.school_name
        """
    ).fetchdf())

    print("\n── Discipline by type (2023-2024) ──")
    print(conn.execute(
        """
        SELECT
            incident_type,
            COUNT(*) AS incident_count
        FROM discipline
        WHERE school_year = '2023-2024'
        GROUP BY incident_type
        ORDER BY incident_count DESC
        """
    ).fetchdf())

    print("\n── Average GPA by school (2023-2024) ──")
    print(conn.execute(
        """
        SELECT
            school_name,
            ROUND(AVG(gpa), 2) AS avg_gpa
        FROM grades
        WHERE school_year = '2023-2024'
        GROUP BY school_name
        ORDER BY avg_gpa DESC
        """
    ).fetchdf())

    conn.close()