"""
File: stratified_sampling.py

Michel Bierlaire
Mon Aug 04 2025, 10:27:32
"""

from dataclasses import dataclass
from typing import Callable

import numpy as np
from pandas import DataFrame, Series


@dataclass
class Population:
    """Data about the full population"""

    data: DataFrame
    choice_variable: str

    @property
    def population_size(self) -> int:
        """Number of individuals in the population"""
        return len(self.data)

    @property
    def choice_set(self) -> list[int]:
        """List of alternatives in the population"""
        return list(self.data[self.choice_variable].unique())


@dataclass
class SampleTuple:
    """Data structure associating a sample with its name"""

    name: str
    sample: DataFrame


@dataclass
class SegmentSample:
    """Characterization of a segment of the population."""

    name: str
    mask_generator: Callable[
        [DataFrame], Series
    ]  # Generates the mask for a given data frame
    target_share: float  # Share of this segment in the sample. H(x)
    target_share_alternatives: (
        dict[int, float] | None
    )  # the share of each alternative in the sample, within the segment H(i | x). If None, same as the population
    population_share: float | None = (
        None  # Share of this segment in the population. W(x)
    )
    population_share_alternatives: dict[int, float] | None = None


@dataclass
class SamplingStrategy:
    """Structure allowing the analyst to characterize a sampling strategy"""

    name: str
    population: Population
    sample_size: int
    segments: list[SegmentSample] | None  # If None, a simple random sample is applied

    @property
    def masks(self) -> list[Series]:
        """Extracts the masks"""
        if self.segments is None:
            return []
        return [
            segment.mask_generator(self.population.data) for segment in self.segments
        ]

    @property
    def column_target_share(self) -> str:
        """Column containing the target shares for exogenous segments H(x)"""
        return f'{self.name}_target_share'

    @property
    def column_population_share(self) -> str:
        """Column containing the population shares for exogenous segments W(x)"""
        return f'{self.name}_population_share'

    @property
    def column_target_share_alt(self) -> str:
        """Prefix of the column containing the target shares per alternative H(i | x)"""
        return f'{self.name}_target_share_alt'

    @property
    def column_population_share_alt(self) -> str:
        """Column containing the population shares per alternative W(i | x)"""
        return f'{self.name}_population_share_alt'

    @property
    def column_exo_probability(self) -> str:
        """Name of the column containing the sampling probability of the exogenous segments"""
        return f'{self.name}_exo_proba'

    @property
    def columns_correction_factor_prefix(self) -> str:
        """Name of the columns containing correction factor for each alternative. The number of each alternative will
        be appended."""
        return f'{self.name}_correction_factor'

    @property
    def column_sampling_probability(self) -> str:
        """Name of the column containing the sampling probability of each observation."""
        return f'{self.name}_sampling_proba'

    @property
    def column_weight(self) -> str:
        """Name of the column with the weights for WESML"""
        return f'{self.name}_weight'

    @property
    def column_random_draw(self) -> str:
        """Name of the column storing the uniform random draws for sampling."""
        return 'random_draws'

    def __post_init__(self) -> None:
        """Performs the validity check just after constructing the object"""
        if not self.check_validity():
            raise ValueError(
                'Invalid SamplingStrategy: The provided values do not pass the validity check.'
            )
        print(f'Validity of {self.name} checked.')

        # Populates the database with the sampling probability.
        self.sampling_probability()

    def mask_statistics(self, dataframe: DataFrame) -> str:
        """Calculates the number of entries corresponding to the masks in a dataframe"""
        if self.segments is None:
            return f'{self.name}: simple random sampling. Number of elements: {len(dataframe)}'
        result = ''
        for segment in self.segments:
            segment_mask = segment.mask_generator(dataframe)
            result += (
                f'*** Segment {segment.name} : {segment_mask.sum()} elements ***\n'
            )
            for alt_id in self.population.choice_set:
                alt_mask = dataframe[self.population.choice_variable] == alt_id
                restricted_mask = segment_mask & alt_mask
                result += f'\tAlt {alt_id} : {restricted_mask.sum()} elements\n'
        return result

    def __str__(self) -> str:
        """print the details of the sampling strategy"""
        if self.segments is None:
            return (
                f'{self.name}: simple random sampling. Expected sample size: {self.sample_size}. Population size: '
                f'{self.population.population_size}. Sampling probability: '
                f'{self.sample_size / self.population.population_size:.3g}'
            )
        result = f'Expected sample size: {self.sample_size}. Population size: {self.population.population_size}\n'

        for segment in self.segments:
            result += f'*** Segment {segment.name} ***\n'
            result += f'Target share H(x):     {segment.target_share:.3g}\n'
            result += f'Population share W(x): {segment.population_share:.3g}\n'
            expected_number_segment = segment.target_share * self.sample_size
            result += f'Expected number in sample: {expected_number_segment:.3g}\n'
            for alt_id in self.population.choice_set:
                if segment.target_share_alternatives is not None:
                    result += (
                        f'\tTarget share alt {alt_id} H({alt_id} | x):    '
                        f'{segment.target_share_alternatives[alt_id]:.3g}\n'
                    )
                    expected_number_alt = (
                        expected_number_segment
                        * segment.target_share_alternatives[alt_id]
                    )
                else:
                    result += (
                        f'\tTarget share alt {alt_id} H({alt_id} | x):    '
                        f'{segment.population_share_alternatives[alt_id]:.3g}\n'
                    )
                    expected_number_alt = (
                        expected_number_segment
                        * segment.population_share_alternatives[alt_id]
                    )
                result += (
                    f'\tPopulation share alt {alt_id} W({alt_id} | x): '
                    f'{segment.population_share_alternatives[alt_id]:.3g}\n'
                )
                result += f'\tExpected number in sample alt {alt_id}: {expected_number_alt:.3g}\n'
        return result

    def is_exogenous(self) -> bool:
        """Determines if the sampling strategy is exogenous"""
        if self.segments is None:
            return True
        for segment in self.segments:
            if segment.target_share_alternatives is not None:
                return False
        return True

    def check_validity(self) -> bool:
        """Check the validity of the sampling strategy:

        - the target shares must sum up to one,
        - the masks must form a partition,
        - when the alternative is defined, it must be consistent with the mask.
        """
        if self.segments is None:
            return True

        valid = True

        # Verify the sum of the target shares.
        total_target_shares = sum(
            [the_segment.target_share for the_segment in self.segments]
        )
        if not np.isclose(total_target_shares, 1.0):
            print(
                f'The target shares should sum up to one. The sum is {total_target_shares}.'
            )
            valid = False

        # Verify that the masks do not intersect
        no_intersection = all(
            (m1 & m2).sum() == 0
            for i, m1 in enumerate(self.masks)
            for m2 in self.masks[i + 1 :]
        )
        if not no_intersection:
            print(f'Some segments overlap.')
            valid = False

        # Verify that the masks collectively form the entire DataFrame
        full_coverage = (sum(self.masks)).sum() == self.population.population_size
        if not full_coverage:
            print('The masks do not cover the full population.')
            valid = False

        # Verify that the alternative is either None or not None for all segments.
        exogenous_sample = self.segments[0].target_share_alternatives is None

        # Verify that all other elements have the same alternative status
        consistent_alternative = all(
            (segment.target_share_alternatives is None) == exogenous_sample
            for segment in self.segments
        )
        if not consistent_alternative:
            print(
                'The target_share_alternatives should all be None, or all be defined.'
            )
            valid = False

        return valid

    def calculate_population_shares(self) -> None:
        """Calculate population shares for each segment"""
        if self.segments is None:
            return
        for segment in self.segments:
            the_mask = segment.mask_generator(self.population.data)
            segment.population_share = the_mask.sum() / self.population.population_size
            self.population.data[self.column_target_share] = segment.target_share
            self.population.data[self.column_population_share] = (
                segment.population_share
            )
            subset = self.population.data[the_mask]
            assert len(subset) == the_mask.sum()
            proportions = subset[self.population.choice_variable].value_counts(
                normalize=True
            )
            segment.population_share_alternatives = {
                alt_id: proportions.get(alt_id, 0)
                for alt_id in self.population.choice_set
            }
            for alt_id in self.population.choice_set:
                population_share = segment.population_share_alternatives.get(alt_id, 0)
                if segment.target_share_alternatives is None:
                    self.population.data[f'{self.column_target_share_alt}_{alt_id}'] = (
                        population_share
                    )
                else:
                    self.population.data[f'{self.column_target_share_alt}_{alt_id}'] = (
                        segment.target_share_alternatives.get(alt_id, 0)
                    )
                self.population.data[f'{self.column_population_share_alt}_{alt_id}'] = (
                    population_share
                )

    def correction_factors(self) -> None:
        """Within each segment, we calculate the conditional probability af being sampled,
        conditional on each alternative"""
        self.calculate_population_shares()
        # Check if they are already calculated
        already_calculated = True
        for alt_id in self.population.choice_set:
            column_name = f'{self.columns_correction_factor_prefix}_{int(alt_id)}'
            if column_name not in self.population.data.columns:
                already_calculated = False
        if already_calculated:
            print('Correction factors have already been calculated. Nothing is done.')
            return

        if self.segments is None:
            # Simple random sample
            for alt_id in self.population.choice_set:
                self.population.data[
                    f'{self.columns_correction_factor_prefix}_{int(alt_id)}'
                ] = (self.sample_size / self.population.population_size)
            return

        for segment in self.segments:
            the_mask = segment.mask_generator(self.population.data)
            for alt_id in self.population.choice_set:
                column_name = f'{self.columns_correction_factor_prefix}_{int(alt_id)}'
                if segment.target_share_alternatives is None:
                    self.population.data.loc[the_mask, column_name] = (
                        segment.target_share
                        * self.sample_size
                        / (segment.population_share * self.population.population_size)
                    )
                else:
                    if (
                        segment.population_share_alternatives.get(alt_id, 0) == 0
                        or segment.population_share_alternatives.get(alt_id, 0) == 0
                    ):
                        self.population.data.loc[the_mask, column_name] = 0.0
                    else:

                        value_to_insert = (
                            segment.target_share
                            * segment.target_share_alternatives.get(alt_id, 0)
                            * self.sample_size
                            / (
                                segment.population_share
                                * segment.population_share_alternatives.get(alt_id, 0)
                                * self.population.population_size
                            )
                        )
                        self.population.data.loc[the_mask, column_name] = (
                            value_to_insert
                        )

    def sampling_probability(self) -> None:
        """Calculate the sampling probability for each observation. It is simply the correction factor corresponding
        to the chosen alternative."""
        self.correction_factors()
        self.population.data[self.column_sampling_probability] = (
            self.population.data.apply(
                lambda row: row[
                    f"{self.columns_correction_factor_prefix}_{int(row[self.population.choice_variable])}"
                ],
                axis=1,
            )
        )
        # Calculate the weights for WESML
        self.population.data[self.column_weight] = self.population.data[
            self.column_sampling_probability
        ].apply(lambda x: 1.0 / x if x != 0.0 else 0.0)

    def sample(self) -> SampleTuple:
        """Extract a sample using the sampling strategy"""

        # Draw a random number for each row
        self.population.data[self.column_random_draw] = np.random.uniform(
            0, 1, size=self.population.population_size
        )

        # Select the rows according to the sampling probability
        sampled_df = self.population.data[
            self.population.data[self.column_random_draw]
            < self.population.data[self.column_sampling_probability]
        ]

        # Normalize the weights so that the sum of the weights is equal to the sample size
        weight_sum = sampled_df[self.column_weight].sum()
        normalization_factor = len(sampled_df) / weight_sum
        sampled_df.loc[:, self.column_weight] = (
            sampled_df[self.column_weight] * normalization_factor
        )
        return SampleTuple(name=self.name, sample=sampled_df)