"""A highly flexible model for forecasting search clicks for a given budget"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/model/00_search_model.ipynb.

# %% auto 0
__all__ = ['ImpressionPrior', 'ImpressionModel']

# %% ../../nbs/model/00_search_model.ipynb 3
from enum import Enum

import pymc as pm
import numpy as np
import pandas as pd
import pytensor.tensor as pt
import xarray as xr

from search_forecast.wrapper.pymc_wrapper import (
    Prior, HSGP, HSGPPeriodic, YearlyFourier, WeeklyFourier,
    Data)
from ..utils.data_gen import hill

# %% ../../nbs/model/00_search_model.ipynb 4
class ImpressionPrior(str,Enum):
    """
    Enum for prior distributions of impressions.
    """
    fraction_search_volume_available = "fraction_search_volume_available"
    competitor_pressure = "competitor_pressure"

# %% ../../nbs/model/00_search_model.ipynb 5
class ImpressionModel:
    """A class to represent the impression model."""
    def __init__(self, 
        dims: tuple[str],
        priors: dict[ImpressionPrior, dict[str, str|float]] | None = None,
        ):
        self.dims = dims
        if priors is None:
            priors = self._default_priors()
        
        self.priors = self._default_priors() | priors
        
        self.model = None
    
    def _process_priors(self):
        """Process the priors for the model."""
        processed_priors = {}
        for prior, params in self.priors.items():
            processed_priors[prior] = Prior(str(prior), **params)
        return processed_priors
        
    def _default_priors(self):
        """Default priors for the model."""
        return {
            ImpressionPrior.fraction_search_volume_available:  dict(prior_name="Beta", alpha=0.5, beta=0.5),
            ImpressionPrior.competitor_pressure: dict(prior_name="HalfCauchy", beta=1)
        }

    def create_model(
        self,
        search_volume: pt.TensorVariable, # Search Volume
        budget: pt.TensorVariable, # Budget
        model: pm.Model | None = None # Model Context,
        ) -> pt.TensorVariable:
        """A model to convert search volume and budget into impressions."""
        model = pm.modelcontext(model)
        self.model = model
        priors = self._process_priors()
        with model:
            multiplier = priors[ImpressionPrior.fraction_search_volume_available].apply()
            competitor_pressure= priors[ImpressionPrior.competitor_pressure].apply()

            impression_rate = pm.Deterministic("impression_rate", multiplier * hill(budget, competitor_pressure, 1.0), dims=self.dims)
            
            impressions = pm.Binomial("impressions", n=search_volume, p=impression_rate, dims=self.dims)
            return impressions

       
    def observe_impressions(
        self,
        impressions: np.ndarray, # Impressions
        coords: dict[str, np.ndarray] | None = None, # Coordinates

    ):
        """Observe the impressions in the model."""
        if self.model is None:
            raise ValueError("Model is not created yet.")
        
        self.model = pm.observe(self.model, {"impressions": impressions})
        if not coords is None:
            with self.model:
                pm.set_data({'impressions': impressions}, coords=coords)
        return self.model
    
