import os
import copy
import json
import sys

import binary_c_python_api

import binarycpython
from binarycpython.utils.grid_options_defaults import grid_options_defaults_dict
from binarycpython.utils.custom_logging_functions import (
    autogen_C_logging_code,
    binary_c_log_code,
    create_and_load_logging_function,
    temp_custom_logging_dir,
)
from binarycpython.utils.functions import get_defaults


# TODO list
# TODO: add functionality to parse cmdline args
# TODO: add functionality to 'on-init' set arguments
# DONE: add functionality to export the arg string.
# DONE: add functionality to export all the options

# TODO: add functionality to return the initial_abundance_hash
# TODO: add functionality to return the isotope_hash
# TODO: add functionality to return the isotope_list
# TODO: add functionality to return the nuclear_mass_hash
# TODO: add functionality to return the nuclear_mass_list
# TODO: add functionality to return the source_list
# TODO: add functionality to return the ensemble_list

# DONE: add functionality to return the evcode_version_string
# Make this function also an API call. Doest seem to get written to a buffer that is stored into a python object. rather its just written to stdout
# DONE: add functionality to return the evcode_args_list
# TODO: add grid generation script

class Population(object):
    def __init__(self):
        """
        Initialisation function of the population class
        """

        self.defaults = get_defaults()

        # Different sections of options
        self.bse_options = (
            {}
        )  # bse_options is just empty. Setting stuff will check against the defaults to see if the input is correct.
        self.grid_options = grid_options_defaults_dict.copy()
        self.custom_options = {}

        # Argline dict
        self.argline_dict = {}

    ###################################################
    # Argument functions
    ###################################################

    # General flow of generating the arguments for the binary_c call:
    # - user provides parameter and value via set (or manually but that is risky)
    # - The parameter names of these input get compared to the parameter names in the self.defaults; with this, we know that its a valid
    # parameter to give to binary_c.
    # - For a single system, the bse_options will be written as a arg line
    # - For a population the bse_options will get copied to a temp_bse_options dict and updated with all the parameters generated by the grid

    # I will NOT create the argument line by fully writing ALL the defaults and overriding user input, that seems not necessary
    # because by using the get_defaults() function we already know for sure which parameter names are valid for the binary_c version
    # And because binary_c uses internal defaults, its not necessary to explicitly pass them.
    # I do however suggest everyone to export the binary_c defaults to a file, so that you know exactly which values were the defaults.

    def set(self, **kwargs):
        """
        Function to set the values of the population. This is the preferred method to set values of functions, as it 
        provides checks on the input. 

        the bse_options will get populated with all the those that have a key that is present in the self.defaults
        the grid_options will get updated with all the those that have a key that is present in the self.grid_options
        
        If neither of above is met; the key and the value get stored in a custom_options dict.
        """

        for key in kwargs.keys():
            # Filter out keys for the bse_options
            if key in self.defaults.keys():
                print("adding: {}={} to BSE_options".format(key, kwargs[key]))
                self.bse_options[key] = kwargs[key]

            # Filter out keys for the grid_options
            elif key in self.grid_options.keys():
                print("adding: {}={} to grid_options".format(key, kwargs[key]))

                self.grid_options[key] = kwargs[key]
            # The of the keys go into a custom_options dict
            else:
                print(
                    "!! Key doesnt match previously known parameter: adding: {}={} to custom_options".format(
                        key, kwargs[key]
                    )
                )
                self.custom_options[key] = kwargs[key]

    def return_argline(self, parameter_dict=None):
        """
        Function to create the string for the arg line from a parameter dict
        """

        if not parameter_dict:
            parameter_dict = self.bse_options

        argline = "binary_c " 
        # TODO: check if that sort actually works
        for param_name in sorted(parameter_dict):
            argline += "{} {} ".format(param_name, parameter_dict[param_name])
        argline = argline.strip()
        return argline

    def generate_population_arglines_file(self, output_file):
        """
        Function to generate a file that contains all the argument lines that would be given to binary_c if the population had been run
        """

        pass

    def add_grid_variable(self, name, longname, valuerange, resolution, spacingfunc, precode, probdist, dphasevol, parameter_name, condition=None):
        """
        Function to add grid variables to the grid_options.

        TODO: Fix this complex function.

        The execution of the grid generation will be through a nested forloop, 
        and will rely heavily on the eval() functionality of python. Which, in terms of safety is very bad, but in terms of flexibility is very good.

        name: 
            name of parameter
            example: name = 'lnm1'            
        longname: 
            Long name of parameter
            example: longname = 'Primary mass'
        range: 
            Range of values to take
            example: range = [log($mmin),log($mmax)]
        resolution: 
            Resolution of the sampled range (amount of samples)
            example: resolution = $resolution->{m1}
        spacingfunction: 
            Function determining how the range is sampled
            example: spacingfunction = "const(log($mmin),log($mmax),$resolution->{m1})"
        precode: 
            # TODO: think of good description.
            example: precode = '$m1=exp($lnm1);'
        probdist: 
            FUnction determining the probability that gets asigned to the sampled parameter
            example: probdist = 'Kroupa2001($m1)*$m1'
        dphasevol: 
            part of the parameter space that the total probability is calculated with
            example: dphasevol = '$dlnm1'
        condition: 
            condition that has to be met in order for the grid generation to continue
            example: condition = '$self->{_grid_options}{binary}==1'
        """

        # Add grid_variable
        grid_variable = {
            "name": name,
            "longname": longname,
            "valuerange": valuerange,
            "resolution": resolution,
            "spacingfunc": spacingfunc,
            "precode": precode,
            "probdist": probdist,
            "dphasevol": dphasevol,
            "parameter_name": parameter_name,
            "condition": condition,
            "grid_variable_number": len(self.grid_options["grid_variables"]),
        }

        # Load it into the grid_options
        self.grid_options["grid_variables"][grid_variable["name"]] = grid_variable

    ###################################################
    # Return functions
    ###################################################

    def return_population_settings(self):
        """
        Function that returns all the options that have been set.

        Can be combined with json to make a nice file. 
        """

        options = {
            "bse_options": self.bse_options,
            "grid_options": self.grid_options,
            "custom_options": self.custom_options,
        }

        return options

    def return_binary_c_version_info(self, parsed=False):
        """
        Function that returns the version information of binary_c
    
        TODO: Put in a nice dict. 
        """

        version_info = binary_c_python_api.return_version_info().strip()

        if parsed:
            version_info = parse_binary_c_version_info(version_info)

        return version_info

    def parse_binary_c_version_info(self):
        """
        Function that parses the output of the version info that binary_c gives. This again is alot of string manipulation. Sensitive to breaking somewhere along the line. 
        """
        # TODO: write parsing function
        pass

    def return_binary_c_defaults(self):
        """
        Function that returns the defaults of the binary_c version that is used.
        """

        return self.defaults

    def return_all_info(self):
        """
        Function that returns all the information about the population and binary_c
        """

        from binarycpython.utils.functions import get_help_all

        population_settings = self.return_population_settings()
        binary_c_defaults = self.return_binary_c_defaults()
        binary_c_version_info = self.return_binary_c_version_info()
        binary_c_help_all_info = get_help_all(print_help=False, return_dict=True)

        all_info = {}
        all_info["population_settings"] = population_settings
        all_info["binary_c_defaults"] = binary_c_defaults
        all_info["binary_c_version_info"] = binary_c_version_info
        all_info["binary_c_help_all"] = binary_c_help_all_info

        return all_info

    def export_all_info(self, use_datadir=False, outfile=None):
        """
        Function that exports the all_info to a json file

        TODO: if any of the values in the dicts here is of a not-serializable form, then we need to change that to a string or something
        so, use a recursive function that goes over the all_info dict and finds those that fit

        TODO: Fix to write things to the directory. which options do which etc
        """

        all_info = self.return_all_info()

        if use_datadir:
            base_name = os.path.splitext(self.custom_options['base_filename'])[0]
            settings_name =  base_name + '_settings.json'

            # Check directory, make if necessary
            os.makedirs(self.custom_options['data_dir'], exist_ok=True)

            settings_fullname = os.path.join(self.custom_options['data_dir'], settings_name)

            # if not outfile.endswith('json'):
            with open(settings_fullname, "w") as f:
                f.write(json.dumps(all_info, indent=4))

        else:
            # if not outfile.endswith('json'):
            with open(outfile, "w") as f:
                f.write(json.dumps(all_info, indent=4))

    def set_custom_logging(self):
        """
        Function/routine to set all the custom logging so that the function memory pointer is known to the grid.
        """

        # C_logging_code gets priority of C_autogen_code
        if self.grid_options["C_auto_logging"]:
            # Generate real logging code
            logging_line = autogen_C_logging_code(self.grid_options["C_auto_logging"])

            # Generate entire shared lib code around logging lines
            custom_logging_code = binary_c_log_code(logging_line)

            # Load memory adress
            self.grid_options[
                "custom_logging_func_memaddr"
            ] = create_and_load_logging_function(custom_logging_code)
        #
        if self.grid_options["C_logging_code"]:
            # Generate entire shared lib code around logging lines
            custom_logging_code = binary_c_log_code(self.grid_options["C_logging_code"])

            # Load memory adress
            self.grid_options[
                "custom_logging_func_memaddr"
            ] = create_and_load_logging_function(custom_logging_code)

    ###################################################
    # Evolution functions
    ###################################################

    def evolve_single(self, parse_function=None):
        """
        Function to run a single system
        
        The output of the run gets returned, unless a parse function is given to this function. 
        """

        ### Custom logging code:
        self.set_custom_logging()

        # Get argument line
        argline = self.return_argline(self.bse_options)

        # Run system
        out = binary_c_python_api.run_system(
            argline,
            self.grid_options["custom_logging_func_memaddr"],
            self.grid_options["store_memaddr"],
        )  # Todo: change this to run_binary again but then build in checks in binary

        # TODO: add call to function that cleans up the temp customlogging dir, and unloads the loaded libraries.

        if parse_function:
            parse_function(self, out)

        else:
            return out

    def evolve_population(self, custom_arg_file=None):
        """
        The function that will evolve the population. This function contains many steps
        """

        ### Custom logging code:
        # C_logging_code gets priority of C_autogen_code
        if self.grid_options["C_auto_logging"]:
            # Generate real logging code
            logging_line = autogen_C_logging_code(self.grid_options["C_auto_logging"])

            # Generate entire shared lib code around logging lines
            custom_logging_code = binary_c_log_code(logging_line)

            # Load memory adress
            self.grid_options[
                "custom_logging_func_memaddr"
            ] = create_and_load_logging_function(custom_logging_code)
        #
        if self.grid_options["C_logging_code"]:
            # Generate entire shared lib code around logging lines
            custom_logging_code = binary_c_log_code(self.grid_options["C_logging_code"])

            # Load memory adress
            self.grid_options[
                "custom_logging_func_memaddr"
            ] = create_and_load_logging_function(custom_logging_code)

        ### Load store
        self.grid_options["store_memaddr"] = binary_c_python_api.return_store("")

        # Execute.
        # TODO: CHange this part alot. This is not finished whatsoever
        out = binary_c_python_api.run_population(
            self.return_argline(),
            self.grid_options["custom_logging_func_memaddr"],
            self.grid_options["store_memaddr"],
        )

        print(out)

        ### Arguments
        # If user inputs a file containing arg lines then use that
        if custom_arg_file:
            # check if file exists
            if os.path.isfile(custom_arg_file):
                # load file
                with open(custom_arg_file) as f:
                    # Load lines into list
                    temp = f.read().splitlines()

                    # Filter out all the lines that dont start with binary_c
                    population_arglines = [
                        line for line in temp if line.startswith("binary_c")
                    ]

        else:
            # generate population from options

            pass

        quit()

        #######
        # Do stuff
        for line in population_arglines:
            print(line)

        pass

        # TODO: add call to function that cleans up the temp customlogging dir, and unloads the loaded libraries.

    ###################################################
    # Testing functions
    ###################################################

    def test_evolve_single(self):
        m1 = 15.0  # Msun
        m2 = 14.0  # Msun
        separation = 0  # 0 = ignored, use period
        orbital_period = 4530.0  # days
        eccentricity = 0.0
        metallicity = 0.02
        max_evolution_time = 15000
        argstring = "binary_c M_1 {0:g} M_2 {1:g} separation {2:g} orbital_period {3:g} eccentricity {4:g} metallicity {5:g} max_evolution_time {6:g}  ".format(
            m1,
            m2,
            separation,
            orbital_period,
            eccentricity,
            metallicity,
            max_evolution_time,
        )

        output = binary_c_python_api.run_binary(argstring)

        print("\n\nBinary_c output:")
        print(output)

    ###################################################
    # Unordered functions
    ###################################################   


    def generate_grid_code(self):
        """
        Function that generates the code from which the population will be made.

        # TODO: make a generator for this.  
        # TODO: Add correct logging everywhere
        # TODO: add different types of grid. 
        """

        code_string = ""
        depth = 0
        indent = '    '

        # Set some values in the generated code:

        # TODO: add imports
        # TODO: 
        code_string += "from binarycpython.utils.probability_distributions import *\n"  
        code_string += "import math\n"
        code_string += "import numpy as np\n"
        code_string += "\n\n"

        code_string += "starcount = 0\n"
        code_string += "probabilities = {}\n"
        code_string += "parameter_dict = {}\n"
        
        # Prepare the probability
        for el in sorted(self.grid_options["grid_variables"].items(), key=lambda x: x[1]['grid_variable_number']):
            grid_variable = el[1]
            code_string += 'probabilities["{}"] = 0\n'.format(grid_variable['parameter_name'])

        # Generate code
        print("Generating grid code")
        for el in sorted(self.grid_options["grid_variables"].items(), key=lambda x: x[1]['grid_variable_number']):
            print("Constructing/adding: {}".format(el[0]))
            grid_variable = el[1]

            # Adding for loop structure
            code_string += indent * depth + 'for {} in {}:'.format(grid_variable['name'], grid_variable['spacingfunc']) + '\n'

            # Add pre-code
            code_string += indent * (depth + 1) + '{}'.format(grid_variable['precode'].replace('\n', '\n'+indent * (depth + 1))) + '\n'

            # Calculate probability
            code_string += indent * (depth + 1) + 'probabilities["{}"] = {}'.format(grid_variable['parameter_name'], grid_variable['probdist']) + '\n'

            # some testing line.
            # code_string += indent * (depth + 1) + 'print({})'.format(grid_variable['name']) + '\n'

            # Add value to dict
            code_string += indent * (depth + 1) + 'parameter_dict["{}"] = {}'.format(grid_variable['parameter_name'], grid_variable['parameter_name']) + '\n'


            # Add some space
            code_string += "\n"

            # increment depth
            depth += 1

        # placeholder for calls to threading
        code_string += indent * (depth) + 'print(parameter_dict)\n'
        # starcount
        code_string += indent * (depth) + 'starcount += 1\n'
        code_string += indent * (depth) + 'print(probabilities)\n'
        code_string += indent * (depth) + 'print("starcount: ", starcount)\n'

        # Write to file
        gridcode_filename = os.path.join(temp_custom_logging_dir(), 'example_grid.py')
        with open(gridcode_filename, 'w') as f:
            f.write(code_string)
################################################################################################