from abc import ABC
from datetime import timedelta
from random import randint, choices

from timeloop import Timeloop

from jesse import sync_publish
from jesse.services.redis import process_status
from multiprocessing import Process, Manager
import click
import numpy as np
import pydash
import jesse.helpers as jh
import jesse.services.logger as logger
from jesse.store import store
import os
from jesse import exceptions
from jesse.modes.optimize_mode.fitness import get_and_add_fitness_to_the_bucket, create_baby
from jesse.routes import router
from numpy import ndarray
from multiprocessing import cpu_count
from jesse.services.progressbar import Progressbar


class Optimizer(ABC):
    def __init__(
            self,
            training_candles: ndarray,
            testing_candles: ndarray,
            optimal_total: int,
            cpu_cores: int,
            csv: bool,
            export_json: bool,
            start_date: str,
            finish_date: str,
            charset: str = r'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvw',
            fitness_goal: float = 1,
    ) -> None:
        if len(router.routes) != 1:
            raise NotImplementedError('optimize_mode mode only supports one route at the moment')

        self.strategy_name = router.routes[0].strategy_name
        self.exchange = router.routes[0].exchange
        self.symbol = router.routes[0].symbol
        self.timeframe = router.routes[0].timeframe
        strategy_class = jh.get_strategy_class(self.strategy_name)
        self.strategy_hp = strategy_class.hyperparameters(None)
        solution_len = len(self.strategy_hp)

        if solution_len == 0:
            raise exceptions.InvalidStrategy('Targeted strategy does not implement a valid hyperparameters() method.')

        self.started_index = 0
        self.start_time = jh.now_to_timestamp()
        self.population = []
        self.iterations = 2000 * solution_len
        self.population_size = solution_len * 100
        self.solution_len = solution_len
        self.charset = charset
        self.fitness_goal = fitness_goal
        self.cpu_cores = 0
        self.optimal_total = optimal_total
        self.training_candles = training_candles
        self.testing_candles = testing_candles
        self.average_execution_seconds = 0

        # check for termination event once per second
        tl_0 = Timeloop()
        @tl_0.job(interval=timedelta(seconds=1))
        def check_for_termination():
            if process_status() != 'started':
                raise exceptions.Termination
        tl_0.start()

        options = {
            'strategy_name': self.strategy_name,
            'exchange': self.exchange,
            'symbol': self.symbol,
            'timeframe': self.timeframe,
            'strategy_hp': self.strategy_hp,
            'csv': csv,
            'json': export_json,
            'start_date': start_date,
            'finish_date': finish_date,
        }

        self.options = {} if options is None else options
        os.makedirs('./storage/temp/optimize', exist_ok=True)
        self.temp_path = f"./storage/temp/optimize/{self.options['strategy_name']}-{self.options['exchange']}-{self.options['symbol']}-{self.options['timeframe']}-{self.options['start_date']}-{self.options['finish_date']}.pickle"

        if fitness_goal > 1 or fitness_goal < 0:
            raise ValueError('fitness scores must be between 0 and 1')

        if not optimal_total > 0:
            raise ValueError('optimal_total must be bigger than 0')

        # # if temp file exists, load data to resume previous session
        # if jh.file_exists(self.temp_path) and click.confirm(
        #         'Previous session detected. Do you want to resume?', default=True
        # ):
        #     self.load_progress()

        if cpu_cores > cpu_count():
            raise ValueError(f'Entered cpu cores number is more than available on this machine which is {cpu_count()}')
        elif cpu_cores == 0:
            self.cpu_cores = cpu_count()
        else:
            self.cpu_cores = cpu_cores

    def generate_initial_population(self) -> None:
        """
        generates the initial population
        """
        length = int(self.population_size / self.cpu_cores)

        progressbar = Progressbar(length)
        for i in range(length):
            people = []
            with Manager() as manager:
                dna_bucket = manager.list([])
                workers = []

                try:
                    for _ in range(self.cpu_cores):
                        dna = ''.join(choices(self.charset, k=self.solution_len))
                        w = Process(
                            target=get_and_add_fitness_to_the_bucket,
                            args=(
                                dna_bucket, jh.get_config('env.optimization'), router.formatted_routes, router.formatted_extra_routes,
                                self.strategy_hp, dna, self.training_candles, self.testing_candles,
                                self.optimal_total
                            )
                        )
                        w.start()
                        workers.append(w)

                    # join workers
                    for w in workers:
                        w.join()
                        if w.exitcode > 0:
                            logger.error(f'a process exited with exitcode: {w.exitcode}')
                except exceptions.Termination:
                    self._handle_termination(manager, workers)

                for d in dna_bucket:
                    people.append({
                        'dna': d[0],
                        'fitness': d[1],
                        'training_log': d[2],
                        'testing_log': d[3]
                    })

            # update dashboard
            self.update_progressbar(progressbar)

            # general_info streams
            general_info = {
                'started_at': jh.timestamp_to_arrow(self.start_time).humanize(),
                'index': f'{(i + 1) * self.cpu_cores}/{self.population_size}',
                'errors_info_count': f'{len(store.logs.errors)}/{len(store.logs.info)}',
                'trading_route': f'{router.routes[0].exchange}, {router.routes[0].symbol}, {router.routes[0].timeframe}, {router.routes[0].strategy_name}',
                'average_execution_seconds': self.average_execution_seconds
            }
            if jh.is_debugging():
                general_info['population_size'] = self.population_size
                general_info['iterations'] = self.iterations
                general_info['solution_length'] = self.solution_len
            sync_publish('general_info', general_info)

            for p in people:
                self.population.append(p)

        sync_publish('progressbar', {
            'current': 100,
            'estimated_remaining_seconds': 0
        })
        # sort the population
        self.population = list(sorted(self.population, key=lambda x: x['fitness'], reverse=True))

    def select_person(self) -> dict:
        # len(self.population) instead of self.population_size because some DNAs might not have been created due to errors
        # to fix an issue with being less than 100 population_len (which means there's only on hyperparameter in the strategy)
        population_len = len(self.population)
        if population_len == 0:
            raise IndexError('population is empty')
        count = int(population_len / 100)
        if count == 0:
            count = 1
        random_index = np.random.choice(population_len, count, replace=False)
        chosen_ones = [self.population[r] for r in random_index]
        return pydash.max_by(chosen_ones, 'fitness')

    def evolve(self) -> list:
        """
        the main method, that runs the evolutionary algorithm
        """
        # clear the logs to start from a clean slate
        jh.clear_file('storage/logs/optimize-mode.txt')

        logger.log_optimize_mode('Optimization session started')

        if self.started_index == 0:
            logger.log_optimize_mode(
                f"Generating {self.population_size} population size (random DNAs) using {self.cpu_cores} CPU cores"
            )
            self.generate_initial_population()

            if len(self.population) < 0.5 * self.population_size:
                msg = f'Too many errors! less than half of the expected population size could be generated. Only {len(self.population)} indviduals from planned {self.population_size} are usable.'
                logger.log_optimize_mode(msg)
                raise ValueError(msg)

            # if even our best individual is too weak, then we better not continue
            if self.population[0]['fitness'] == 0.0001:
                msg = 'Cannot continue because no individual with the minimum fitness-score was found. Your strategy seems to be flawed or maybe it requires modifications. '
                logger.log_optimize_mode(msg)
                raise exceptions.InvalidStrategy(msg)

        loop_length = int(self.iterations / self.cpu_cores)

        i = self.started_index
        progressbar = Progressbar(loop_length)
        while i < loop_length:
            with Manager() as manager:
                people_bucket = manager.list([])
                workers = []

                try:
                    for _ in range(self.cpu_cores):
                        mommy = self.select_person()
                        daddy = self.select_person()
                        w = Process(
                            target=create_baby,
                            args=(
                                people_bucket, mommy, daddy, self.solution_len, self.charset,
                                jh.get_config('env.optimization'), router.formatted_routes,
                                router.formatted_extra_routes,
                                self.strategy_hp, self.training_candles, self.testing_candles,
                                self.optimal_total
                            )
                        )
                        w.start()
                        workers.append(w)

                    for w in workers:
                        w.join()
                        if w.exitcode > 0:
                            logger.error(f'a process exited with exitcode: {w.exitcode}')
                except exceptions.Termination:
                    self._handle_termination(manager, workers)

                # update dashboard
                click.clear()
                self.update_progressbar(progressbar)
                # general_info streams
                general_info = {
                    'started_at': jh.timestamp_to_arrow(self.start_time).humanize(),
                    'index': f'{(i + 1) * self.cpu_cores}/{self.iterations}',
                    'errors_info_count': f'{len(store.logs.errors)}/{len(store.logs.info)}',
                    'trading_route': f'{router.routes[0].exchange}, {router.routes[0].symbol}, {router.routes[0].timeframe}, {router.routes[0].strategy_name}',
                    'average_execution_seconds': self.average_execution_seconds
                }
                if jh.is_debugging():
                    general_info['population_size'] = self.population_size
                    general_info['iterations'] = self.iterations
                    general_info['solution_length'] = self.solution_len
                sync_publish('general_info', general_info)

                if self.population_size > 50:
                    number_of_ind_to_show = 40
                elif self.population_size > 20:
                    number_of_ind_to_show = 15
                elif self.population_size > 9:
                    number_of_ind_to_show = 9
                else:
                    raise ValueError('self.population_size cannot be less than 10')

                best_candidates = [{
                        'rank': j + 1,
                        'dna': self.population[j]['dna'],
                        'fitness': round(self.population[j]['fitness'], 4),
                        'training_win_rate': self.population[j]['training_log']['win-rate'],
                        'training_total_trades': self.population[j]['training_log']['total'],
                        'training_pnl': self.population[j]['training_log']['PNL'],
                        'testing_win_rate': self.population[j]['testing_log']['win-rate'],
                        'testing_total_trades': self.population[j]['testing_log']['total'],
                        'testing_pnl': self.population[j]['testing_log']['PNL'],
                    } for j in range(number_of_ind_to_show)]
                sync_publish('best_candidates', best_candidates)

                # one person has to die and be replaced with the newborn baby
                for baby in people_bucket:
                    # never kill our best performer
                    random_index = randint(1, len(self.population) - 1)
                    self.population[random_index] = baby
                    self.population = list(sorted(self.population, key=lambda x: x['fitness'], reverse=True))

                    # reaching the fitness goal could also end the process
                    if baby['fitness'] >= self.fitness_goal:
                        self.update_progressbar(progressbar, finished=True)
                        sync_publish('alert', {
                            'message': f'Fitness goal reached after iteration {i*self.cpu_cores}',
                            'type': 'success'
                        })
                        return baby

                # TODO: bring back progress resumption
                # # save progress after every n iterations
                # if i != 0 and int(i * self.cpu_cores) % 50 == 0:
                #     self.save_progress(i)

                # TODO: bring back
                # # store a take_snapshot of the fittest individuals of the population
                # if i != 0 and i % int(100 / self.cpu_cores) == 0:
                #     self.take_snapshot(i * self.cpu_cores)

                i += 1

        sync_publish('alert', {
            'message': f"Finished {self.iterations} iterations. Check your best DNA candidates, "
                       f"if you don't like any of them, try modifying your strategy.",
            'type': 'success'
        })
        return self.population

    def run(self) -> list:
        return self.evolve()


    @staticmethod
    def _handle_termination(manager, workers):
        logger.info('Terminating session...')

        # terminate all workers
        for w in workers:
            w.terminate()

        # shutdown the manager process manually since garbage collection cannot won't get to do it for us
        manager.shutdown()

        # now we can terminate the main session safely
        raise exceptions.Termination

    def update_progressbar(self, progressbar, finished=False):
        if finished:
            progressbar.finish()
        else:
            progressbar.update()
        self.average_execution_seconds = progressbar.average_execution_seconds / self.cpu_cores
        sync_publish('progressbar', {
            'current': progressbar.current,
            'estimated_remaining_seconds': progressbar.estimated_remaining_seconds
        })


    # def save_progress(self, iterations_index: int) -> None:
    #     """
    #     pickles data so we can later resume optimizing
    #     """
    #     data = {
    #         'population': self.population,
    #         'iterations': self.iterations,
    #         'iterations_index': iterations_index,
    #         'population_size': self.population_size,
    #         'solution_len': self.solution_len,
    #         'charset': self.charset,
    #         'fitness_goal': self.fitness_goal,
    #         'options': self.options,
    #     }
    #
    #     with open(self.temp_path, 'wb') as f:
    #         pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

    # def load_progress(self) -> None:
    #     """
    #     Unpickles data to resume from previous optimizing session population
    #     """
    #     with open(self.temp_path, 'rb') as f:
    #         data = pickle.load(f)
    #
    #     self.population = data['population']
    #     self.iterations = data['iterations']
    #     self.started_index = data['iterations_index']
    #     self.population_size = data['population_size']
    #     self.solution_len = data['solution_len']
    #     self.charset = data['charset']
    #     self.fitness_goal = data['fitness_goal']
    #     self.options = data['options']

    # def take_snapshot(self, index: int) -> None:
    #     """
    #     stores a snapshot of the fittest population members into a file.
    #     """
    #     study_name = f"{self.options['strategy_name']}-{self.options['exchange']}-{self.options['symbol']}-{self.options['timeframe']}-{self.options['start_date']}-{self.options['finish_date']}"
    #
    #     dna_json = {'snapshot': []}
    #     for i in range(30):
    #         dna_json['snapshot'].append(
    #             {'iteration': index, 'dna': self.population[i]['dna'], 'fitness': self.population[i]['fitness'],
    #              'training_log': self.population[i]['training_log'], 'testing_log': self.population[i]['testing_log'],
    #              'parameters': jh.dna_to_hp(self.options['strategy_hp'], self.population[i]['dna'])})
    #
    #     path = f'./storage/genetics/{study_name}.txt'
    #     os.makedirs('./storage/genetics', exist_ok=True)
    #     txt = ''
    #     with open(path, 'a', encoding="utf-8") as f:
    #         txt += '\n\n'
    #         txt += f'# iteration {index}'
    #         txt += '\n'
    #
    #         for i in range(30):
    #             log = f"win-rate: {self.population[i]['training_log']['win-rate']} %, total: {self.population[i]['training_log']['total']}, PNL: {self.population[i]['training_log']['PNL']} % || win-rate: {self.population[i]['testing_log']['win-rate']} %, total: {self.population[i]['testing_log']['total']}, PNL: {self.population[i]['testing_log']['PNL']} %"
    #
    #             txt += '\n'
    #             txt += f"{i + 1} ==  {self.population[i]['dna']}  ==  {self.population[i]['fitness']}  ==  {log}"
    #
    #         f.write(txt)
    #
    #     if self.options['csv']:
    #         path = f'storage/genetics/csv/{study_name}.csv'
    #         os.makedirs('./storage/genetics/csv', exist_ok=True)
    #         exists = os.path.exists(path)
    #
    #         df = json_normalize(dna_json['snapshot'])
    #
    #         with open(path, 'a', newline='', encoding="utf-8") as outfile:
    #             if not exists:
    #                 # header of CSV file
    #                 df.to_csv(outfile, header=True, index=False, encoding='utf-8')
    #
    #             df.to_csv(outfile, header=False, index=False, encoding='utf-8')
    #
    #     if self.options['json']:
    #         path = f'storage/genetics/json/{study_name}.json'
    #         os.makedirs('./storage/genetics/json', exist_ok=True)
    #         exists = os.path.exists(path)
    #
    #         mode = 'r+' if exists else 'w'
    #         with open(path, mode, encoding="utf-8") as file:
    #             if not exists:
    #                 snapshots = {"snapshots": []}
    #                 snapshots["snapshots"].append(dna_json['snapshot'])
    #                 json.dump(snapshots, file, ensure_ascii=False)
    #             else:
    #                 # file exists - append
    #                 file.seek(0)
    #                 data = json.load(file)
    #                 data["snapshots"].append(dna_json['snapshot'])
    #                 file.seek(0)
    #                 json.dump(data, file, ensure_ascii=False)
    #             file.write('\n')
