src/population/variant_generator.py

"""
Persona variant generator for population sampling

Creates realistic variations of base personas while maintaining core identity.
"""

import random
from typing import List
from enum import Enum
from copy import deepcopy

from ..personas.models import Persona, Demographics, Psychographics


class VariationLevel(str, Enum):
    """Level of variation to apply to persona variants"""
    CONSERVATIVE = "conservative"  # Small variations (±1-2 points)
    MODERATE = "moderate"  # Medium variations (±2-3 points)
    DIVERSE = "diverse"  # Large variations (±3-5 points)


class PersonaVariant:
    """A variant of a base persona with statistical variations"""

    def __init__(self, base_persona: Persona, variant_id: int, variation_level: VariationLevel):
        """
        Create a variant of a base persona

        Args:
            base_persona: The base persona to create a variant from
            variant_id: Unique identifier for this variant (0-99 for population of 100)
            variation_level: How much variation to apply
        """
        self.base_persona = base_persona
        self.variant_id = variant_id
        self.variation_level = variation_level

        # Create the variant persona
        self.persona = self._create_variant()

    def _create_variant(self) -> Persona:
        """Create a variant persona with applied variations"""
        # Deep copy the base persona
        variant = deepcopy(self.base_persona)

        # Update the persona ID
        variant.persona_id = f"{self.base_persona.persona_id}_v{self.variant_id}"

        # Get variation ranges based on level
        ranges = self._get_variation_ranges()

        # Apply variations to demographics
        variant.demographics = self._vary_demographics(
            variant.demographics, ranges
        )

        # Apply variations to psychographics
        variant.psychographics = self._vary_psychographics(
            variant.psychographics, ranges
        )

        return variant

    def _get_variation_ranges(self) -> dict:
        """Get variation ranges based on variation level"""
        if self.variation_level == VariationLevel.CONSERVATIVE:
            return {
                "age": 2,
                "years_in_community": 1,
                "scalar_traits": 1,  # For 1-10 scales
            }
        elif self.variation_level == VariationLevel.MODERATE:
            return {
                "age": 5,
                "years_in_community": 3,
                "scalar_traits": 2,
            }
        else:  # DIVERSE
            return {
                "age": 10,
                "years_in_community": 5,
                "scalar_traits": 3,
            }

    def _vary_demographics(self, demographics: Demographics, ranges: dict) -> Demographics:
        """Apply variations to demographic attributes"""
        varied = deepcopy(demographics)

        # Vary age within reasonable bounds
        age_delta = random.randint(-ranges["age"], ranges["age"])
        varied.age = max(18, min(100, demographics.age + age_delta))

        # Vary years in community
        years_delta = random.randint(-ranges["years_in_community"], ranges["years_in_community"])
        varied.years_in_community = max(0, demographics.years_in_community + years_delta)

        # Occasionally vary household size (20% chance)
        if random.random() < 0.2:
            varied.household_size = max(1, demographics.household_size + random.choice([-1, 1]))

        # Occasionally flip boolean traits (10% chance each)
        if random.random() < 0.1:
            varied.has_children = not demographics.has_children
        if random.random() < 0.1:
            varied.owns_home = not demographics.owns_home

        return varied

    def _vary_psychographics(self, psychographics: Psychographics, ranges: dict) -> Psychographics:
        """Apply variations to psychographic attributes"""
        varied = deepcopy(psychographics)

        # Vary scalar traits (all the 1-10 scales)
        trait_range = ranges["scalar_traits"]

        varied.openness_to_change = self._vary_scalar(
            psychographics.openness_to_change, trait_range
        )
        varied.community_engagement = self._vary_scalar(
            psychographics.community_engagement, trait_range
        )
        varied.environmental_concern = self._vary_scalar(
            psychographics.environmental_concern, trait_range
        )
        varied.economic_focus = self._vary_scalar(
            psychographics.economic_focus, trait_range
        )
        varied.social_equity_focus = self._vary_scalar(
            psychographics.social_equity_focus, trait_range
        )
        varied.risk_tolerance = self._vary_scalar(
            psychographics.risk_tolerance, trait_range
        )

        # Occasionally shift political leaning (15% chance)
        if random.random() < 0.15:
            varied.political_leaning = self._shift_political_leaning(
                psychographics.political_leaning
            )

        return varied

    def _vary_scalar(self, value: int, range_val: int) -> int:
        """Vary a scalar value (1-10 scale) within range"""
        delta = random.randint(-range_val, range_val)
        return max(1, min(10, value + delta))

    def _shift_political_leaning(self, current: str) -> str:
        """Shift political leaning by one step"""
        leanings = [
            "very_progressive",
            "progressive",
            "moderate",
            "conservative",
            "very_conservative"
        ]

        if current not in leanings:
            return current

        idx = leanings.index(current)
        # Shift left or right by one (50/50 chance)
        direction = random.choice([-1, 1])
        new_idx = max(0, min(len(leanings) - 1, idx + direction))

        return leanings[new_idx]


class VariantGenerator:
    """Generator for creating persona variants"""

    def __init__(self, base_persona: Persona, variation_level: VariationLevel):
        """
        Initialize variant generator

        Args:
            base_persona: The base persona to create variants from
            variation_level: How much variation to apply
        """
        self.base_persona = base_persona
        self.variation_level = variation_level
        self._counter = 0

    def generate_variant(self, suffix: str = None) -> Persona:
        """
        Generate a single variant

        Args:
            suffix: Optional suffix for variant ID

        Returns:
            A variant Persona object
        """
        variant_id = self._counter if suffix is None else suffix
        self._counter += 1
        variant_obj = PersonaVariant(self.base_persona, variant_id, self.variation_level)
        return variant_obj.persona


def generate_variants(
    base_persona: Persona,
    population_size: int,
    variation_level: VariationLevel,
    seed: int = None
) -> List[Persona]:
    """
    Generate a population of persona variants

    Args:
        base_persona: The base persona to create variants from
        population_size: Number of variants to generate
        variation_level: How much variation to apply
        seed: Random seed for reproducibility (optional)

    Returns:
        List of Persona objects representing the population
    """
    if seed is not None:
        random.seed(seed)

    variants = []
    for i in range(population_size):
        variant_obj = PersonaVariant(base_persona, i, variation_level)
        variants.append(variant_obj.persona)

    return variants