hr_env/server/data_gen.py

"""Synthetic company generation using faker and numpy distributions."""

from __future__ import annotations

import random
from typing import Optional

import numpy as np
from faker import Faker

from hr_env.server.company import DEPARTMENT_NAMES, Company, Department
from hr_env.server.employee import Employee

# Department-specific role templates
ROLES = {
    "Engineering": [
        ("Software Engineer", 1), ("Software Engineer", 2), ("Senior Engineer", 3),
        ("Staff Engineer", 4), ("Principal Engineer", 5),
    ],
    "Sales": [
        ("Sales Rep", 1), ("Account Executive", 2), ("Senior AE", 3),
        ("Sales Manager", 4), ("VP Sales", 5),
    ],
    "Operations": [
        ("Operations Analyst", 1), ("Operations Specialist", 2), ("Operations Manager", 3),
        ("Senior Ops Manager", 4), ("VP Operations", 5),
    ],
    "HR": [
        ("HR Coordinator", 1), ("HR Specialist", 2), ("HR Manager", 3),
        ("Senior HR Manager", 4), ("VP People", 5),
    ],
    "Finance": [
        ("Financial Analyst", 1), ("Senior Analyst", 2), ("Finance Manager", 3),
        ("Controller", 4), ("CFO", 5),
    ],
}

# Salary base by level (lognormal parameters)
SALARY_BASE = {1: 55000, 2: 75000, 3: 100000, 4: 135000, 5: 180000}
SALARY_SIGMA = 0.15  # lognormal spread

# Department size distribution (fraction of total)
DEPT_SIZE_FRACTIONS = {
    "Engineering": 0.35,
    "Sales": 0.25,
    "Operations": 0.20,
    "HR": 0.08,
    "Finance": 0.12,
}

# Skills by department
DEPT_SKILLS = {
    "Engineering": ["python", "java", "cloud", "ml", "devops", "frontend", "backend", "data"],
    "Sales": ["negotiation", "crm", "pipeline", "presentation", "territory", "closing"],
    "Operations": ["logistics", "process", "lean", "supply_chain", "quality", "planning"],
    "HR": ["recruiting", "compliance", "training", "benefits", "employee_relations"],
    "Finance": ["accounting", "budgeting", "forecasting", "audit", "tax", "reporting"],
}


def generate_company(
    seed: int = 42,
    size: int = 300,
    name: str = "Simulated Corp",
    base_revenue: Optional[float] = None,
    hr_budget: Optional[float] = None,
) -> Company:
    """Generate a synthetic company with realistic employee distributions.

    Args:
        seed: Random seed for reproducibility.
        size: Total number of employees (200-500 recommended).
        name: Company name.
        base_revenue: Annual baseline revenue. Defaults to size * 150_000.
        hr_budget: Annual HR budget. Defaults to size * 6000.
    """
    rng = np.random.default_rng(seed)
    fake = Faker()
    Faker.seed(seed)

    if base_revenue is None:
        base_revenue = size * 150_000
    if hr_budget is None:
        hr_budget = size * 6000

    company = Company(
        name=name,
        base_revenue=base_revenue,
        hr_budget=hr_budget,
        hr_budget_remaining=hr_budget / 4,  # First quarter allocation
        rng=random.Random(seed),  # seeded stream for deterministic turnover
    )

    for dept_name in DEPARTMENT_NAMES:
        dept_size = max(5, int(size * DEPT_SIZE_FRACTIONS[dept_name]))
        dept = Department(name=dept_name)
        roles = ROLES[dept_name]
        skills_pool = DEPT_SKILLS[dept_name]

        for _ in range(dept_size):
            # Level distribution: pyramid shape
            level_probs = np.array([0.35, 0.30, 0.20, 0.10, 0.05])
            level = int(rng.choice([1, 2, 3, 4, 5], p=level_probs))

            # Find matching role
            role_name = next((r for r, l in roles if l == level), roles[0][0])

            # Salary: lognormal around base for level
            base = SALARY_BASE[level]
            salary = float(rng.lognormal(np.log(base), SALARY_SIGMA))
            salary = round(max(35000, salary), -2)  # Round to nearest 100

            # Performance: normal(3.2, 0.8) clipped [1, 5]
            perf = float(np.clip(rng.normal(3.2, 0.8), 1.0, 5.0))

            # Tenure: exponential(lambda=24 months)
            tenure = int(rng.exponential(24))

            # Engagement: beta(7, 3) * 100
            engagement = float(rng.beta(7, 3) * 100)

            # Skills: 2-4 random from department pool
            n_skills = min(len(skills_pool), int(rng.integers(2, 5)))
            skills = list(rng.choice(skills_pool, size=n_skills, replace=False))

            emp = Employee(
                id=f"{dept_name[:3].lower()}_{fake.unique.random_int(min=1000, max=9999)}",
                name=fake.name(),
                department=dept_name,
                role=role_name,
                level=level,
                salary=salary,
                tenure_months=tenure,
                performance_score=perf,
                engagement=engagement,
                skills=skills,
                training_hours=float(rng.exponential(10)),
            )

            # Derived attributes
            emp.promotability = float(np.clip(
                0.1 + perf * 0.12 + (tenure / 60) * 0.1 + rng.normal(0, 0.1), 0, 1
            ))
            emp.transferability = float(np.clip(
                0.3 + len(skills) * 0.08 + rng.normal(0, 0.1), 0, 1
            ))
            emp.retainability = float(np.clip(
                0.4 + engagement / 200 + (tenure / 48) * 0.1 + rng.normal(0, 0.1), 0, 1
            ))
            emp.update_flight_risk()

            dept.employees.append(emp)

        company.departments[dept_name] = dept

    return company