"""
Generates synthetic trees where the root molecule optimizes for a specific objective
based on Therapeutic Data Commons (TDC) oracle functions. Uses a genetic algorithm
to optimize embeddings before decoding.
"""


import os
import numpy as np 
from tqdm import tqdm 
import sys
path_here = os.path.dirname(os.path.realpath(__file__))
sys.path.append(path_here)
sys.path.append('.')
from main.optimizer import BaseOptimizer, Objdict

from syn_net.utils.ga_utils import crossover, mutation
import multiprocessing as mp
import numpy as np
import pandas as pd
import time
import scripts._mp_decode as decode
from syn_net.utils.predict_utils import mol_fp
from tqdm import tqdm 


def fitness(embs, _pool, obj):
    """
    Returns the scores for the root molecules in synthetic trees generated by the
    input molecular embeddings.

    Args:
        embs (list): Contains molecular embeddings (vectors).
        _pool (mp.Pool): A pool object, which represents a pool of workers (used
            for multiprocessing).
        obj (str): The objective function to use to compute the fitness.

    Raises:
        ValueError: Raised if the specified objective function is not implemented.

    Returns:
        scores (list): Contains the scores for the root molecules in the
            generated trees.
        smiles (list): Contains the root molecules encoded as SMILES strings.
        trees (list): Contains the synthetic trees generated from the input
            embeddings.
    """
    results = _pool.map(decode.func, embs)
    print("fitness decode finished")
    ### debug mode: without pool (parallel)
    # results = []
    # for emb in embs:
    #     results.append(decode.func(emb))
    ### debug mode 

    smiles  = [r[0] for r in results]
    trees   = [r[1] for r in results]
    scores = [obj(smi) for smi in smiles]
    return scores, smiles, trees


def distribution_schedule(n, total):
    """
    Determines the type of probability to use in the `crossover` function, based
    on the number of generations which have occured.

    Args:
        n (int): Number of elapsed generations.
        total (int): Total number of expected generations.

    Returns:
        str: Describes a type of probability distribution.
    """
    if n < 4 * total/5:
        return 'linear'
    else:
        return 'softmax_linear'

def num_mut_per_ele_scheduler(n, total):
    """
    Determines the number of bits to mutate in each vector, based on the number
    of elapsed generations.

    Args:
        n (int): Number of elapsed generations.
        total (int): Total number of expected generations.

    Returns:
        int: Number of bits to mutate.
    """
    return 24

def mut_probability_scheduler(n, total):
    """
    Determines the probability of mutating a vector, based on the number of elapsed
    generations.

    Args:
        n (int): Number of elapsed generations.
        total (int): Total number of expected generations.

    Returns:
        float: The probability of mutation.
    """
    if n < total/2:
        return 0.5
    else:
        return 0.5



class SynNet_Optimizer(BaseOptimizer):

    def __init__(self, args=None):
        super().__init__(args)
        self.model_name = "synnet"

    def _optimize(self, oracle, config):

        self.oracle.assign_evaluator(oracle)
        config = Objdict(config)

        if config['restart']:
            population = np.load(config.input_file)
            print(f"Starting with {len(population)} fps from {config.input_file}")
        else:
            try:
                starting_smiles = pd.read_csv(config.input_file).sample(config.num_population)
                starting_smiles = starting_smiles['smiles'].tolist()
                population = np.array(
                    [mol_fp(smi, config['radius'], config['nbits']) for smi in starting_smiles]
                )
                # print(population.shape, 'population shape ------')
                # population = population.reshape((population.shape[0], population.shape[2]))
                print(f"Starting with {len(starting_smiles)} fps from {config.input_file}")
            except:
                starting_smiles = self.all_smiles[:config.num_population]
                population = np.array(
                    [mol_fp(smi, config.radius, config.nbits) for smi in starting_smiles]
                )
                # population = np.ceil(np.random.random(size=(config['num_population'], config['nbits'])) * 2 - 1)
                print(f"Starting with {config.num_population} fps with {config.nbits} bits")

        with mp.Pool(processes=config.ncpu) as pool:
            scores, mols, trees = fitness(embs=population,
                                          _pool=pool,
                                          obj=self.oracle)
        scores     = np.array(scores)
        score_x    = np.argsort(scores)
        population = population[score_x[::-1]]
        mols       = [mols[i] for i in score_x[::-1]]
        scores     = scores[score_x[::-1]]
        print(f"Initial: {scores.mean():.3f} +/- {scores.std():.3f}")

        recent_scores = []
        patience = 0 

        dist_            = "softmax_linear"
        num_mut_per_ele_ = config.num_mut_per_ele
        mut_probability_ = config.mut_probability

        n = 0
        patience = 0

        while True:

            n += 1

            if len(self.oracle) > 100:
                self.sort_buffer()
                old_score = np.mean([item[1][0] for item in list(self.mol_buffer.items())[:100]])
            else:
                old_score = 0

            t = time.time()

            # dist_            = distribution_schedule(n, config.num_gen)
            # num_mut_per_ele_ = num_mut_per_ele_scheduler(n, config.num_gen)
            # mut_probability_ = mut_probability_scheduler(n, config.num_gen)

            offspring = crossover(parents=population,
                                  offspring_size=config.num_offspring,
                                  distribution=dist_)
            offspring = mutation(offspring_crossover=offspring,
                                 num_mut_per_ele=num_mut_per_ele_,
                                 mut_probability=mut_probability_)

            new_population = np.unique(np.concatenate([population, offspring], axis=0), axis=0)

            with mp.Pool(processes=config.ncpu) as pool:
                new_scores, new_mols, trees = fitness(new_population, pool, self.oracle)

            if self.finish:
                print('max oracle hit, abort ...... ')
                break 

            new_scores = np.array(new_scores)
            scores = []
            mols = []

            parent_idx = 0
            indices_to_print = []
            while parent_idx < config.num_population:
                max_score_idx = np.where(new_scores == np.max(new_scores))[0][0]
                if new_mols[max_score_idx] not in mols:
                    indices_to_print.append(max_score_idx)
                    scores.append(new_scores[max_score_idx])
                    mols.append(new_mols[max_score_idx])
                    population[parent_idx, :] = new_population[max_score_idx, :]
                    new_scores[max_score_idx] = -999999
                    parent_idx += 1
                else:
                    new_scores[max_score_idx] = -999999

            scores = np.array(scores)
            print(f"Generation {n+1}: {scores.mean():.3f} +/- {scores.std():.3f}")
            print(f"Consumed time: {(time.time() - t):.3f} s")

            recent_scores.append(scores.mean())
            if len(recent_scores) > 10:
                del recent_scores[0]

            if len(self.oracle) > 100:
                self.sort_buffer()
                new_score = np.mean([item[1][0] for item in list(self.mol_buffer.items())[:100]])
                if (new_score - old_score) < 1e-3:
                    patience += 1
                    if patience >= self.args.patience:
                        self.log_intermediate(finish=True)
                        print('convergence criteria met, abort ...... ')
                        break
                else:
                    patience = 0

                old_score = new_score
