data_generator.py

"""
data_generator.py
Generates synthetic but realistic financial snapshots
for Easy, Medium, and Hard tasks.
"""

import random
from typing import TypedDict, List, Dict


class FinancialSnapshot(TypedDict):
    year: int
    net_income: float
    shareholder_equity: float
    total_debt: float
    revenue: float
    gross_profit: float


# ------------------------------------------------------------
# Core generator
# ------------------------------------------------------------

def _single_year(rng: random.Random, base_rev: float) -> FinancialSnapshot:
    revenue = round(base_rev * rng.uniform(0.8, 1.2), 2)
    gross_profit = round(revenue * rng.uniform(0.05, 0.35), 2)
    net_income = round(gross_profit * rng.uniform(-0.2, 1.0), 2)
    equity = round(rng.uniform(50, 800), 2)
    debt = round(equity * rng.uniform(0.2, 3.5), 2)

    return {
        "year": 0,
        "net_income": net_income,
        "shareholder_equity": equity,
        "total_debt": debt,
        "revenue": revenue,
        "gross_profit": gross_profit,
    }


# ------------------------------------------------------------
# Required by Medium task
# ------------------------------------------------------------

def generate_snapshot(seed: int | None = None) -> FinancialSnapshot:
    rng = random.Random(seed)
    snap = _single_year(rng, rng.uniform(200, 1500))
    snap["year"] = 2024
    return snap


# ------------------------------------------------------------
# Used by Easy task
# ------------------------------------------------------------

def generate_single_year(seed: int | None = None) -> FinancialSnapshot:
    return generate_snapshot(seed)


# ------------------------------------------------------------
# Used by Hard task
# ------------------------------------------------------------

def generate_three_years(seed: int | None = None) -> List[FinancialSnapshot]:
    rng = random.Random(seed)
    base = rng.uniform(300, 1500)

    snaps = []
    for yr in [2022, 2023, 2024]:
        s = _single_year(rng, base)
        s["year"] = yr
        base *= rng.uniform(0.85, 1.15)
        snaps.append(s)

    return snaps


# ------------------------------------------------------------
# Ratio computation (used by Medium/Hard graders)
# ------------------------------------------------------------

def compute_ratios(s: FinancialSnapshot) -> Dict:
    roe = round(s["net_income"] / s["shareholder_equity"], 4)
    dte = round(s["total_debt"] / s["shareholder_equity"], 4)
    pm = round(s["gross_profit"] / s["revenue"], 4)
    return {
        "roe": roe,
        "debt_to_equity": dte,
        "profit_margin": pm,
    }

# ------------------------------------------------------------
# Health classification (required by Medium grader)
# ------------------------------------------------------------

def classify_health(roe: float, dte: float, pm: float) -> str:
    """
    Classifies financial health based on ratio thresholds.
    Matches logic expected by grader_medium.
    """

    if roe > 0.15 and dte < 1.0 and pm > 0.2:
        return "Healthy"
    elif roe > 0.05 and dte < 2.0 and pm > 0.1:
        return "At Risk"
    else:
        return "Critical"