from pathlib import Path
from re import sub

template_codes = {
    'simulation' : lambda class_name: f'''
"""
Simulation template
"""

import numpy as np
from harissa.core.simulation import Simulation

class {class_name}(Simulation):
    """
    {class_name} simulation.
    """
    def __init__(self):
        """
        This is a dummy init. Feel free to change its implementation or 
        to delete it.
        """ 
        self.verbose = False

    def run(self, time_points, initial_state, parameter):
        """
        Perform a constant simulation. Feel free to change its implementation.

        Parameters
        ----------
        time_points: np.ndarray
            Recorded time points.

        initial_state: np.ndarray
            initial rna and protein levels. 
            The shape of this array is (2, G) 
            with G the number of genes (including the stimulus).
            The first row is the rna level and 
            the second row is the protein level at the start of the simulation.
            The first column of the protein levels (initial_state[1, 0]) 
            represents the stimulus. 

        parameter: NetworkParameter
            network parameter used to simulate (here not used).
            The network parameter contains the following attributes:

            - Bursting parameters:
                - burst_frequency_min # k[0]
                - burst_frequency_max # k[1]
                - burst_size_inv      # b

            - Creation rates:
                - creation_rna        # s[0]
                - creation_protein    # s[1]

            - Degradation rates:
                - degradation_rna     # d[0]
                - degradation_protein # d[1]

            - Network parameters:
                - basal               # beta
                - interaction         # theta

            See https://harissa-framework.github.io/harissa/main/_autogenerated/harissa.core.parameter.html#module-harissa.core.parameter
            for more informations on this object.

        Returns
        -------
        Simulation.Result
            Data class to store the time points, rna levels and protein levels.
            See https://harissa-framework.github.io/harissa/main/_autogenerated/harissa.core.simulation.html
            for more informations.
        """
        # The shape of the rna and protein levels are (T, G) 
        # with T the number of time points
        # and G the number of genes (including the stimulus)
        # The axis 0 represents the time points id,
        # the axis 1 represents the genes level. 
        shape = (time_points.size, initial_state.shape[1])
        
        # create a 2D array that will record 
        # the rna level for every time points
        rna_levels = np.zeros(shape)

        # copy the initial rna level for every time_points
        rna_levels[:] = initial_state[0]

        # create a 2D array that will record 
        # the protein level for every time points
        protein_levels = np.zeros(shape)
        
        # copy the initial protein level for every time_points
        protein_levels[:] = initial_state[1]

        # Store the time points, rna levels and protein levels
        # inside a Simulation.Result object
        res = self.Result(time_points, rna_levels, protein_levels)
        
        # If you need to add other attributes, 
        # you can dynamically define them later 
        # or create a subclass of Simulation.Result.
        # For example:
        # res.my_variable = my_value or setattr(res, "my_variable", my_value)

        return res

'''.lstrip(),
    'inference': lambda class_name: f'''
"""
Inference template
"""

import numpy as np
from harissa.core.parameter import NetworkParameter
from harissa.core.inference import Inference
from harissa.core.dataset import Dataset

class {class_name}(Inference):
    """
    {class_name} inference.
    """
    def __init__(self):
        """
        This is a dummy init. Feel free to change its implementation or 
        to delete it.
        """ 
        self.verbose = False

    @property
    def directed(self):
        """
        True if the {class_name} infers an 
        asymmetric interaction matrix else False.

        Returns
        -------
        bool
        """
        return True

    def run(self, data, parameter):
        """
        Perform an identity inference. Feel free to change its implementation.

        Parameters
        ----------
        data: Dataset
            dataset used to infer (here not used)
            The dataset contains the following attributes:
            
            - time_points
            - count_matrix
            - gene_names

            See https://harissa-framework.github.io/harissa/main/_autogenerated/harissa.core.dataset.html
            for more informations.

        parameter: NetworkParameter
            network parameter updated after inferring.
            The network parameter contains the following attributes:

            - Bursting parameters:
                - burst_frequency_min # k[0]
                - burst_frequency_max # k[1]
                - burst_size_inv      # b

            - Creation rates:
                - creation_rna        # s[0]
                - creation_protein    # s[1]

            - Degradation rates:
                - degradation_rna     # d[0]
                - degradation_protein # d[1]

            - Network parameters:
                - basal               # beta
                - interaction         # theta

            See https://harissa-framework.github.io/harissa/main/_autogenerated/harissa.core.parameter.html#module-harissa.core.parameter
            for more informations on this object.

        Returns
        -------
        Inference.Result
            Wrapper class around the updated network and other custom data.

        """
        return self.Result(parameter)

'''.lstrip()
}

def template(args):
    if args.path.absolute() == Path.cwd().absolute():
        path = args.path / args.path.absolute().name
    else:
        path = args.path

    path = path.with_suffix('.py') 
    code = template_codes[args.template_type](
        sub(r'[_-]+', ' ', path.stem).title().replace(' ', '')
    )

    with open(path, 'w') as python_file:
        python_file.write(code)

def add_subcommand(main_subparsers):
    parser = main_subparsers.add_parser(
        'template',
        help='generate a template for Inference or Simulation subclass'
    )
    parser.add_argument(
        'template_type',
        choices=tuple(template_codes.keys()),
        help='Type of template that will be generated.'
    )
    parser.add_argument(
        'path', 
        type=Path,
        help='path where to generate the template. '
             'The name of the template is deducted from it.'
    )

    parser.set_defaults(run=template)
