from __future__ import annotations

import argparse
import functools
import pathlib
import pickle
import random
from typing import Any

import pandas as pd
from dynpric import simulate_market
from dynpric.demand.ferreira2018 import BernoulliDemand
from dynpric.firms.ferreira2018 import sample_demand
from dynpric.firms.ferreira2018 import SamplingStrategies
from dynpric.priors import BetaPrior
from dynpric.types import Belief
from dynpric.types import DemandRealized
from dynpric.types import Firm
from dynpric.types import History
from dynpric.types import PriceLevel
from dynpric.types import PricesSet


DATA_DIR = pathlib.Path(__file__).parent / 'data'


COMMANDS = {}

def command(fn):
    COMMANDS[fn.__name__] = fn
    return fn

def logger(
    firm: Firm,
    prices_set: PricesSet,
    demand: DemandRealized,
) -> dict[str, Any] | None:

    return {
        'price': prices_set[firm],
        'demand': demand[firm],
        'revenue': prices_set[firm] * demand[firm],
        **{
            str(belief.price) + '_alpha': belief.prior.α
            for belief in firm.beliefs
        },
        **{
            str(belief.price) + '_beta': belief.prior.β
            for belief in firm.beliefs
        },
    }


def beliefs():
    return [
        Belief(29.9, BetaPrior(1, 1)),
        Belief(34.9, BetaPrior(1, 1)),
        Belief(39.9, BetaPrior(1, 1)),
        Belief(44.9, BetaPrior(1, 1)),
    ]


def _observe_market(self, history: History) -> None:
    last_period = history[-1]
    price_set = last_period.prices[self]
    demand = last_period.demand[self]
    # Update belief
    (belief,) = [belief for belief in self.beliefs if belief.price == price_set]
    belief.prior.update(int(demand))


class ThompsonSeller:
    def __init__(self, name, beliefs):
        self.name = name
        self.beliefs = beliefs

    @property
    def price(self):
        demand = [
            sample_demand(belief.prior, SamplingStrategies.thompson)
            for belief in self.beliefs
        ]
        prices = [belief.price for belief in self.beliefs]

        revenue = [q * p for q, p in zip(demand, prices)]
        idx = revenue.index(max(revenue))
        return prices[idx]

    def observe_market(self, history: History) -> None:
        _observe_market(self, history)


class GreedySeller:
    def __init__(self, name, beliefs):
        self.name = name
        self.beliefs = beliefs

    @property
    def price(self):
        demand = [
            sample_demand(belief.prior, SamplingStrategies.greedy)
            for belief in self.beliefs
        ]
        prices = [belief.price for belief in self.beliefs]

        revenue = [q * p for q, p in zip(demand, prices)]
        idx = revenue.index(max(revenue))
        return prices[idx]

    def observe_market(self, history: History) -> None:
        _observe_market(self, history)


class εGreedySeller:
    def __init__(self, name, beliefs, ε):
        self.name = name
        self.beliefs = beliefs
        self.ε = ε

    @property
    def price(self):
        demand = [
            sample_demand(belief.prior, SamplingStrategies.greedy)
            for belief in self.beliefs
        ]
        prices = [belief.price for belief in self.beliefs]
        if random.random() > self.ε:
            # Greedy
            revenue = [q * p for q, p in zip(demand, prices)]
            idx = revenue.index(max(revenue))
            return prices[idx]
        else:
            # Explore
            return random.choice(prices)

    def observe_market(self, history: History) -> None:
        _observe_market(self, history)


def run_trial(firm, trial_id, n_periods):

    if trial_id % 50 == 0:
        print(f'Trial id {trial_id}')

    demand = BernoulliDemand(
        price_levels=[
            PriceLevel(29.9, 0.8),
            PriceLevel(34.9, 0.6),
            PriceLevel(39.9, 0.3),
            PriceLevel(44.9, 0.1),
        ]
    )

    return simulate_market(
        n_periods=n_periods,
        firms=[firm],
        demand=demand,
        logger=logger,
    )


def main():
    args = parse_args()
    return COMMANDS[args.command]()


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('command')
    return parser.parse_args()


@command
def simulate():
    N_PERIODS = 500
    N_TRIALS = 3000

    SELLERS = {
        'εgreedy': lambda: εGreedySeller('εGreedy', beliefs(), 0.05),
        'thompson': lambda: ThompsonSeller('thompson', beliefs()),
        'greedy': lambda: GreedySeller('greedy', beliefs()),
    }

    for name, seller in SELLERS.items():
        print('Simulating', name)
        seller_trial_factory = functools.partial(run_trial, seller)
        results = [
            run_trial(seller(), trial_id, N_PERIODS)
            for trial_id in range(N_TRIALS)
        ]

        with open(DATA_DIR / f'{name}.pickle', 'wb') as f:
            pickle.dump(results, f)


@command
def generate_csv():
    def process_trial(trial, trial_id):
        logs, history = trial
        return [
            process_period(period, trial_id, period_id)
            for period_id, period in enumerate(logs)
        ]

    def process_period(period, trial_id, period_id):
        # There's only one firm
        ((firm, info),) = period.items()
        return {'trial_id': trial_id, 'period_id': period_id, **info}

    def process_pickled_results(file):
        with open(file, 'rb') as f:
            results = pickle.load(f)

        df = pd.concat(
            [
                pd.DataFrame(process_trial(trial, trial_id))
                for trial_id, trial in enumerate(results)
            ]
        )

        df.to_csv(DATA_DIR / f'{file.stem}.csv', index=False)

    for file in DATA_DIR.glob('*.pickle'):
        process_pickled_results(file)




if __name__ == '__main__':
    exit(main())
