"""Markov Chain Monte Carlo for trDesign."""

# native
from datetime import datetime
import time

# lib
import numpy as np
from numpy.random import choice
import torch
from torch.cuda.amp import autocast
import math
import random
import time

# pkg
from losses import *  # pylint: disable=wildcard-import, unused-wildcard-import
from tr_Rosetta_model import trRosettaEnsemble, prep_seq
from utils import aa2idx, distogram_distribution_to_distogram, idx2aa, plot_progress
import config as cfg
from utils import definegroupbydist, plot_muts

def v(torch_value):
    """Return a detached value, if possible."""
    try:
        return torch_value.cpu().detach().item()
    except Exception:
        return torch_value

def motifsort(elem):
    return elem[4]

def placemotifs(motifs, seq_L, sequence, mode = 0):
    """Randomly position discontinous motifs and check if valid. motif = [start, end, length, restraints, group, newstart, newend]"""
    print("Placing motifs...")
    valid = False
    i = 0
    sum = 0

    for m in motifs:
        sum += m[2]
    if (sum > seq_L):
        print("Sequence too short")
        return placemotifs(motifs[:-1], seq_L, sequence, mode)
    elif sum == seq_L:
        pos = 0
        mode = -3
        for m in motifs:
            m[5] = pos
            m[6] = m[5] + m[2] - 1
            pos = pos + m[2]

    while not valid:
        if mode == 0 or mode == 1:
            for m in motifs[:]:
                m[5] = np.random.randint(0,seq_L-m[2]+1)
                m[6] = m[5]+m[2]-1
        elif mode == -2:
            print("predefined motifs")
            seq_con = ""
            for i in range(0,seq_L):
                restraint = "-"
                for m in motifs:
                    if i >= m[5] and i <= m[6]:
                        mi = i - m[5] #local index
                        c = m[3][mi]  #con type
                        if c in cfg.sequence_restraint_letters:
                            restraint = sequence[i]
                        continue

                seq_con = seq_con + restraint

            return motifs, seq_con
        elif mode == -3:
            break;
        elif mode == 2:
            spacing = (seq_L-sum)/(1+len(motifs))
            pos = int(spacing)
            for m in motifs[:]:
                    buffer = int(abs(np.random.normal(spacing,spacing)))
                    m[5] = pos
                    m[6] = m[5]+m[2]-1
                    pos = m[6] + buffer

        elif mode == 2.1:
            spacing = int((seq_L-sum-6)/(len(motifs)-1)) #calculate spacing for even placement with short unrestrained tails
            pos = int(seq_L-sum-((len(motifs)-1)*spacing)) #center the motifs
            for m in motifs[:]:
                    m[5] = pos
                    m[6] = m[5]+m[2]-1
                    pos = m[6]+spacing
        elif mode == 2.2:
            motifs[0][5] = 0
            motifs[0][6] = motifs[0][2]-1
            motifs[-1][6] = seq_L - 1
            motifs[-1][5] = seq_L - motifs[-1][2]
            last = motifs[-1]
            spacing = int((seq_L-sum)/(len(motifs)-1))
            pos = motifs[0][6]
            for m in motifs[1:-1]:
                buffer = int(abs(np.random.normal(spacing,spacing)))
                pos += 1 + buffer
                m[5] = pos
                m[6] = m[5] + m[2] - 1
                pos = m[6]

        elif mode == 3:
            motifs.sort(key=motifsort)
            for m in motifs[:]: #set all same group
                m[4] = 0
            return placemotifs(motifs, seq_L, sequence, mode = 2)
        elif mode == 4 or mode == 5:
            motifs = definegroupbydist(motifs,cfg.target_motif_path, mode)
            return placemotifs(motifs, seq_L, sequence, mode = 3)
        else: return placemotifs(motifs, seq_L, sequence, mode = 1)

        #check if motifs are valid
        valid = True
        i = i + 1
        for m1 in motifs[:]:
            if m1[6] > seq_L - 1:
                print("ends outside")
                valid = False
                continue
            for m2 in motifs[:]:
                if m1 == m2: continue
                if (m1[5] >= m2[5]) and (m1[5] <= m2[6]):
                    valid = False
                    break
                elif m1[6] >= m2[5] and m1[6] <= m2[6]:
                    valid = False
                    break

        if i > 20000: #if not valid, place motifs manually
            print("No valid motif placements found, attempting sequential positions")
            random.shuffle(motifs)
            rest = seq_L - sum
            buffer = math.floor(rest/(len(motifs)+1))
            pos = buffer
            for m in motifs:
                m[5] = pos
                m[6] = m[5] + m[2]-1
                pos += m[2] + buffer
            break

    seq_con = ""

    constrain_seq = False
    if sequence is not None:
        for i in range(0,seq_L):
            restraint = "-"
            for m in motifs:
                if i >= m[5] and i <= m[6]:
                    mi = i - m[5] #local index
                    c = m[3][mi]  #con type
                    if c in cfg.sequence_restraint_letters:
                        si = m[0]+mi  #template index
                        restraint = sequence[si]
                        constrain_seq = True
                    continue

            if i == 0 and cfg.first_residue_met:
                restraint = "M"

            seq_con = seq_con + restraint

    if not constrain_seq:
        seq_con = None

    return motifs, seq_con

def createmask(motifs, seq_L, save_dir, is_site_mask = False):
    """Create mask for discontinous motifs"""

    #setup mask size
    motif_mask_g = np.zeros((seq_L, seq_L))
    motif_mask = np.zeros((seq_L, seq_L))

    for m1 in motifs[:]:
        for i in range(m1[5],m1[6]+1):
            c1 = m1[3][i-m1[5]]
            if c1 in cfg.structure_restraint_letters:  #contraint is structural?
                for m2 in motifs[:]:
                    if math.fabs(m2[4]-m1[4]) > 1: #motifs are not in restrained groups?
                        continue
                    for j in range(m2[5],m2[6]+1):
                        c2 = m2[3][j-m2[5]]
                        if c2 in cfg.structure_restraint_letters:
                            v1 = cfg.structure_restraint_mask_values[c1]
                            v2 = cfg.structure_restraint_mask_values[c2]
                            value = np.clip(v1*v2, 1, 100)
                            motif_mask[i,j] = value
                            motif_mask_g[i,j] = value**0.5

    if not is_site_mask:
        plot_values = motif_mask_g.copy()
        plot_distogram(
            plot_values,
            save_dir / f"motif_mask_groups.jpg",
        )

    return motif_mask

class MCMC_Optimizer(torch.nn.Module):
    """Markov Chain Monte Carlo optimizer."""

    # We don't define `.forward()`, but it's nn-based.
    # pylint: disable=too-many-instance-attributes, abstract-method

    DEFAULT_ROOT_PATH = Path(__file__).parent.parent
    DEFAULT_MODEL_PATH = DEFAULT_ROOT_PATH / "models" / "trRosetta_models"
    DEFAULT_BGND_PATH = DEFAULT_ROOT_PATH / "backgrounds"
    DEFAULT_RESULTS_PATH = DEFAULT_ROOT_PATH / "results"

    def __init__(
        self,
        L,
        aa_weight,
        MCMC,
        native_frequencies,
        experiment_name,
        aa_valid,
        max_aa_index=20,
        sequence_constraint=None,
        target_motif_path=None,
        trRosetta_model_dir="models/trRosetta_models",
        background_distribution_dir="backgrounds",
        motifs=None,
        motifmode = 0,
        motif_weight = 1,
        bkg_weight = 1
    ):
        """Construct the optimizer."""
        super().__init__()
        self.results_dir = self.setup_results_dir(experiment_name)
        self.bkg_dir = background_distribution_dir
        self.structure_models = trRosettaEnsemble(
            trRosetta_model_dir
        )  # .share_memory()
        print(
            f"{self.structure_models.n_models} structure prediction models loaded to {d()}"
        )

        # General params:
        self.eps = 1e-7
        self.seq_L = L
        self.motifs = None
        self.motifmode = motifmode
        self.use_sites = cfg.use_sites
        print("Sequence Length: " + str(self.seq_L))
        if motifs is not None:
            _motifs,_seq_con = placemotifs(motifs, self.seq_L, sequence_constraint, mode=self.motifmode)
            self.motifs = _motifs
            if _seq_con is not None: sequence_constraint = _seq_con

        for m in self.motifs: print(m)
        print(sequence_constraint)

        # Setup MCMC params:
        self.beta, self.N, self.coef, self.M, self.MaxM = (
            MCMC["BETA_START"],
            MCMC["N_STEPS"],
            MCMC["COEF"],
            MCMC["M"],
            MCMC["MAX"],
        )
        self.aa_weight = aa_weight

        # Setup sequence constraints:
        self.aa_valid = aa_valid
        self.native_frequencies = native_frequencies

        self.seq_constraint = sequence_constraint
        if self.seq_constraint is not None:
            assert len(self.seq_constraint) == self.seq_L, \
            "Constraint length (%d) must == Seq_L (%d)" %(len(self.seq_constraint), self.seq_L)

            self.seq_constraint = (
                aa2idx(self.seq_constraint).copy().reshape([1, self.seq_L])
            )
            self.seq_constraint_indices = np.where(
                self.seq_constraint != max_aa_index, 1, 0
            )

        self.target_motif_path = target_motif_path
        self.setup_losses()

        # stats
        self.bad_accepts = []
        self.n_accepted_mutations = 0
        self.n_accepted_bad_mutations = 0
        self.best_metrics = {}
        self.best_step = 0
        self.best_sequence = None
        self.best_E = None
        self.step = 0
        self.motif_weight = motif_weight
        self.bkg_weight = bkg_weight

        self.matrix_setup()

    def matrix_setup(self):
        self.substitution_matrix = {}
        with open("src/" + cfg.MATRIXFILE) as reader:
            columns = reader.readline().strip().split()
            lines = reader.readlines()
            for aa in columns:
                if aa in cfg.RM_AA: continue
                if aa not in cfg.ALPHABET_core_str: continue
                self.substitution_matrix[aa] = {}

            for line in lines:
                data = line.strip().split()
                aa1 = data[0]
                if aa1 not in cfg.ALPHABET_core_str: continue
                sub = {}
                for i in range(0,len(columns)):
                    aa2 = columns[i]
                    if aa2 in cfg.RM_AA: continue
                    if aa2 not in cfg.ALPHABET_core_str: continue
                    value = data[i+1]
                    try: sub[aa2] = float(value)
                    except: sub[aa2] = 0

                self.substitution_matrix[aa1] = dict(sorted(sub.items(), key=lambda x: x[1]))

        if True: #convert to non negative values
            min_value = 1

            for a in self.substitution_matrix.items():
                for b in a[1].items():
                    if b[1] < min_value: min_value = b[1]

            if min_value < 0:
                offset = -min_value + 1

                for a in self.substitution_matrix.items():
                    for b in a[1].items():
                        self.substitution_matrix[a[0]][b[0]] += offset

                #for a in self.substitution_matrix.items():
                #    for b in a[1].items():
                #        self.substitution_matrix[a[0]][b[0]] = self.substitution_matrix[a[0]][b[0]]*self.substitution_matrix[a[0]][b[0]]

    def setup_results_dir(self, experiment_name):
        """Create the directories for the results."""
        results_dir = (
            self.DEFAULT_RESULTS_PATH
            / experiment_name
            / datetime.now().strftime("%Y-%m-%d_%H%M%S")
        )
        results_dir.mkdir(parents=True, exist_ok=True)
        (results_dir / "distogram_evolution").mkdir(parents=True, exist_ok=True)
        print(f"Writing results to {results_dir}")
        return results_dir

    def setup_losses(self):
        """Prepare the loss functions."""

        # Initialize protein background distributions:
        if cfg.BACKGROUND: self.bkg_loss = Structural_Background_Loss(self.seq_L, self.bkg_dir)
        self.aa_bkgr_distribution = torch.from_numpy(self.native_frequencies).to(d())

        # Motif-Loss:
        self.motif_weight = 1.00 #Multiplier for the loss contribution
        self.motif_mask = np.zeros((self.seq_L, self.seq_L)) #LxL tensor, int64, L currently specified in config.py. Cotains only zeros
        self.motif_mask = torch.from_numpy(self.motif_mask).long().to(d())

        #If motif target is specified. Motif target file is .npz LxL. L is overall protein length
        #Overrides previous motif_mask
        #New mask will be ones with a diagonal zero line
        if self.target_motif_path is not None:
            self.motif_mask = np.ones((self.seq_L, self.seq_L))

            if self.motifs is not None:
                self.motif_mask = createmask(self.motifs, self.seq_L, self.results_dir)

            self.motif_mask = torch.from_numpy(self.motif_mask).long().to(d())
            self.motif_mask.fill_diagonal_(0)
            self.motif_sat_loss = Motif_Satisfaction(
                self.target_motif_path, mask=self.motif_mask, save_dir=self.results_dir, motifs = self.motifs
            )

            self.site_weight = 0
            if self.use_sites:
                self.site_mask = createmask(self.motifs, self.seq_L, self.results_dir, is_site_mask = True)
                self.site_mask = torch.from_numpy(self.site_mask).long().to(d())
                self.site_mask.fill_diagonal_(0)

                self.site_sat_loss = Site_Satisfaction(
                    self.target_motif_path, mask=self.site_mask, motifs = self.motifs
                )

        # Apply the background KL-loss only under the hallucination_mask == 1 region
        self.hallucination_mask = 1 - torch.clip(self.motif_mask, 0, 1)
        self.hallucination_mask.fill_diagonal_(0)

    def loss(self, sequence, structure_predictions, msa1hot, track=False):
        """Compute the loss function."""

        # Top-prob: extremily expensive operation
        #TM_score_proxy = top_prob(structure_predictions['dist'], verbose=False)
        #TM_score_proxy = TM_score_proxy[0]  # We're running with batch_size = 1

        # Background KL-loss:
        if cfg.BACKGROUND:
            background_loss = self.bkg_loss(
                structure_predictions, hallucination_mask=self.hallucination_mask
            )
        else:
            background_loss = torch.tensor(0)

        # aa composition loss
        aa_samp = (
            msa1hot[0, :, :20].sum(axis=0) / self.seq_L + self.eps #take entire first sequence, all 1hot chars, and sum across length?
        )  # Get relative frequency for each AA
        aa_samp = (
            aa_samp / aa_samp.sum()
        )  # Normalize to turn into distributions (possibly redundant)
        loss_aa = (
            aa_samp
            * torch.log(aa_samp / (self.aa_bkgr_distribution + self.eps) + self.eps)
        ).sum()

        # Motif Loss:
        if self.target_motif_path is not None: #check if target motif has been specified
            motif_loss, motif_loss_pos = self.motif_sat_loss(structure_predictions) #motif_sat_loss = MotifSatisfaction object (nn object)
        else:
            motif_loss = 0 #no target motif = no loss

        if self.use_sites:
            site_loss = self.site_sat_loss(structure_predictions)
        else:
            site_loss = torch.tensor(0)

        # total loss
        loss_v = (
            self.bkg_weight * background_loss + self.aa_weight * loss_aa + self.motif_weight * motif_loss + self.site_weight * site_loss
        )

        metrics = {}
        if track:
            metrics["aa_weight"] = self.aa_weight
            metrics["background_loss"] = background_loss
            metrics["motif_loss"] = motif_loss
            metrics["site_loss"] = site_loss
            metrics["total_loss"] = loss_v
            #metrics["TM_score_proxy"] = TM_score_proxy

        return loss_v, metrics

    def metropolis(self, seq, seq_curr, E_curr, E):
        """Compute the Metropolis criterion."""
        accepted = True
        deltaE = E_curr - E

        # Metropolis criterion
        if E_curr < E:  # Lower energy, replace!
            seq = np.copy(seq_curr)
            E = E_curr
        else:  # Higher energy, maybe replace..
            if torch.exp((E - E_curr) * self.beta) > np.random.uniform():
                seq = np.copy(seq_curr)
                E = E_curr
                self.bad_accepts.append(1)
                self.n_accepted_bad_mutations += 1
            else:
                accepted = False
                self.bad_accepts.append(0)

        self.register_mutation_fitness(accepted, deltaE)

        # Update the best sequence:
        if E_curr < self.best_E:
            self.best_E = E_curr
            self.best_sequence = idx2aa(seq_curr[0])
            self.best_step = self.step

        return seq, E

    def mutate(self, seq):
        """Return a mutated version of the sequence."""
        seq_curr = np.copy(seq)

        # Introduce a random mutation using the allowed aa_types:
        if cfg.GRADIENT: idx = random.choices(range(self.seq_L), self.gradient, k=1)[0]
        else: idx = np.random.randint(self.seq_L)

        from_aa = idx2aa(seq_curr[0])[idx]

        #Perform mutation
        mutation_list = self.mutation_score[idx]

        if cfg.MATRIX:
            if len(mutation_list) == 0: mutation_list.append([from_aa, from_aa, 0, True, 0]) #Adds accepted dummy mutation

            #mut = self.mutation_score[idx]
            #mut_score = mut[2]
            #accepted = mut[3]
            #mins = min(self.mutation_score, key=lambda x: x[2])[2]
            #maxs = max(self.mutation_score, key=lambda x: x[2])[2]

            ################################
            ### SUBSTITUTION MATRIX MODE ###
            if cfg.MATRIX_MODE == 'probability':
                last_accepted = None
                verybad = ""

                for m in reversed(mutation_list):
                    if m[2] > 0.1:
                        verybad += m[1]
                    if m[3]:
                        last_accepted = m
                        break

                if last_accepted[2] < 0.005: p = 1
                else: p = 0

                for i in range(10):
                    list_of_candidates = [k for k,v in self.substitution_matrix[last_accepted[p]].items()]
                    probability_distribution = [v for k,v in self.substitution_matrix[last_accepted[p]].items()]
                    mutation = random.choices(list_of_candidates, probability_distribution, k=1)[0]
                    if mutation not in verybad: break

                seq_curr[0, idx] = aa2idx(mutation)

            #############################################
            ### AMINO ACID PROPERTY GROUP MATRIX MODE ###
            elif cfg.MATRIX_MODE == 'groups':
                recent_mutations = []

                for m in reversed(mutation_list):
                    if self.step - m[4] > self.seq_L * 10: break
                    recent_mutations.append(m)

                if len(recent_mutations) == 0:
                    seq_curr[0, idx] = np.random.choice(self.aa_valid)
                else:
                    group_scores = [[0,0,None]] * len(cfg.AA_PROPERTY_GROUPS)
                    i = 0
                    for k,v in cfg.AA_PROPERTY_GROUPS.items():
                        group_scores[i][2] = v
                        for m in recent_mutations:
                            if m[1] in v:
                                group_scores[i][1] += m[2]
                                group_scores[i][0] += 1
                        i += 1

                    best = ""
                    best_s = 1000

                    for i in range(len(group_scores)):
                        g = group_scores[i]
                        avg = g[1]/(g[0]+1)
                        if avg < best_s:
                            best = g[2]
                            best_s = avg
                        elif avg == best_s:
                            best += g[2]

                    options = []
                    for aa in best:
                        if aa in cfg.RM_AA: continue
                        options.append(aa)

                    seq_curr[0, idx] = aa2idx(np.random.choice(options))

            ################################
            ### MSA MUTATION MATRIX MODE ###
            elif cfg.MATRIX_MODE == 'msa':
                group = cfg.TEMPLATE_AA_PROPERTIES[idx]
                seq_curr[0, idx] = aa2idx(np.random.choice(group))

            ###########################
            ### MODE NOT RECOGNIZED ###
            else: seq_curr[0, idx] = np.random.choice(self.aa_valid)

        #######################
        ### NON MATRIX MODE ###
        else: seq_curr[0, idx] = np.random.choice(self.aa_valid)

        to_aa = idx2aa(seq_curr[0])[idx]

        if self.seq_constraint is not None:  # Fix the constraint:
            seq_curr = np.where(
                self.seq_constraint_indices, self.seq_constraint, seq_curr
            )

        if np.equal(seq_curr, seq).all(): # If the mutation did not change anything, retry
            self.reverse_log = False
            return self.mutate(seq)

        # Store mutation information
        self.current_mutations.append([idx, from_aa, to_aa])

        return seq_curr

    def select_mutation_options(self, matrix, threshold = 0, above = True):
        if above:
            return [k for k,v in matrix.items() if v >= threshold]
        else:
            return [k for k,v in matrix.items() if v <= threshold]

    def diff_to_weight(self, diff, increase = True):
        if increase:
            if diff == 0: return 1.25
            else: return 1.5 - 0.025*diff
        else: return 0.8 + 0.02 * diff

    def register_mutation_fitness(self, accepted, deltaE):
        if deltaE < -1: deltaE = 0
        else: deltaE = float(deltaE.cpu().detach().numpy())

        good = deltaE < 0

        for mut in self.current_mutations:
            self.good_mutations[mut[0]] += 1

            if good:
                self.good_accepts.append(1)
                for i in range(mut[0]-20,mut[0]+21):
                    if i >= 0 and i < self.seq_L:
                        self.gradient[i] *= self.diff_to_weight(abs(mut[0]-i))
            else:
                self.good_accepts.append(0)
                for i in range(mut[0]-10,mut[0]+11):
                    if i >= 0 and i < self.seq_L:
                        self.gradient[i] *= self.diff_to_weight(abs(mut[0]-i), increase = False)

            self.mutation_log.append([mut[0], mut[1], mut[2], self.substitution_matrix[mut[1]][mut[2]], deltaE, accepted])

            if accepted: self.n_accepted_mutations += 1

            self.mutation_score[mut[0]].append([mut[1], mut[2], deltaE, accepted, self.step]) #from, to, result, accepted, when

            if self.reverse_log:
                print("result:  " + str(mut) + " " + str(deltaE) + " | " + str(accepted))

        self.gradient = np.clip([i+0.01*(1-i) for i in self.gradient], 0.2, 20)
        self.current_mutations = []
        self.reverse_log = False

    def fixup_MCMC(self, seq):
        """Dynamically adjust the metropolis beta parameter to improve performance."""
        if self.step - self.best_step > min(self.N // 4,1000):
            # No improvement for a long time, reload the best_sequence and decrease beta:
            print("reload best seq")
            self.best_step = self.step
            self.beta = self.beta / (self.coef ** 2)
            seq = torch.from_numpy(
                aa2idx(self.best_sequence).reshape([1, self.seq_L])
            ).long()
            self.mutation_score = [[]] * self.seq_L

        elif np.mean(self.bad_accepts[-100:]) < 0.02:
            # There has been some progress recently, but we're no longer accepting any bad mutations...
            self.beta = self.beta / self.coef
        else:
            self.beta = self.beta * self.coef

        self.beta = np.clip(self.beta, 5, self.MaxM)

        #FIX DYNAMIC MATRIX
        if cfg.MATRIX_DYNAMIC: cfg.MATRIX_MODE = random.choice(['msa', 'probability', 'groups', 'none'])

        return seq

    @torch.no_grad()
    def run(self, start_seq):
        """Run the MCMC loop."""
        # pylint: disable=too-many-locals

        start_time = time.time()
        muttime = 0
        nntime = 0
        losstime = 0
        scoretime = 0
        misctime = 0

        # initialize with given input sequence
        print("Initial seq: ", start_seq)
        seq = aa2idx(start_seq).copy().reshape([1, self.seq_L])

        nsave = min(max(1, self.N // 20),50)
        E, E_tracker = np.inf, []
        self.bad_accepts = []
        self.good_accepts = []
        self.n_accepted_mutations = 0
        self.n_accepted_bad_mutations = 0
        self.best_metrics = {}
        self.best_step = 0
        self.best_sequence = start_seq
        self.best_E = E
        self.best_bkg = 0
        self.best_mtf = 0
        self.best_distogram_distribution = None
        self.gradient = [1] * self.seq_L
        self.good_mutations = [0] * self.seq_L
        self.mutation_score = [[]] * self.seq_L
        self.reverse_log = False
        self.current_mutations = []
        self.mutation_log = []
        self.last_update = time.time()
        self.last_graph = time.time()

        # Main loop:
        for self.step in range(self.N + 1):

            # random mutation at random position, also fix sequence constraint
            mut_start = time.time()
            if self.step != 0:
                seq_curr = self.mutate(seq)
                #if cfg.MATRIX_MODE == 'msa':
                #    for i in range(10): seq_curr = self.mutate(seq_curr)
            else: seq_curr = seq

            # Preprocess the sequence, relatively expensive operations (0.04s per pass)
            seq_curr = torch.from_numpy(seq_curr).long()
            model_input, msa1hot = prep_seq(seq_curr)
            muttime += time.time() - mut_start

            # probe effect of mutation
            nn_start = time.time()
            structure_predictions = self.structure_models( #run trRosettaEnsemble -> runs trrosetta n times
                model_input, use_n_models=cfg.n_models
            )
            nntime += time.time() - nn_start

            #loss functions
            loss_start = time.time()
            E_curr, metrics = self.loss(
                seq_curr, structure_predictions, msa1hot, track=True
            )
            losstime += time.time() - loss_start

            #scoring evaluation
            score_start = time.time()
            if E_curr < self.best_E:
                self.best_bkg = metrics["background_loss"]
                self.best_mtf = metrics["motif_loss"]
                self.best_site = metrics["site_loss"]
                self.best_distogram_distribution = structure_predictions['dist']

            seq, E = self.metropolis(seq, seq_curr, E_curr, E)

            E_tracker.append(v(E))
            delta_step_best = self.step - self.best_step
            scoretime += time.time() - score_start

            misc_start = time.time()
            if time.time() - self.last_update > cfg.report_interval or self.step == 0:
                fps = self.step / (time.time() - start_time)
                background_loss = metrics["background_loss"].cpu().detach().numpy()
                mtf_loss = metrics["motif_loss"].cpu().detach().numpy()
                site_loss = metrics["site_loss"].cpu().detach().numpy()

                b_bkg = self.best_bkg.cpu().detach().numpy()
                b_mtf = self.best_mtf.cpu().detach().numpy()
                b_site = self.best_site.cpu().detach().numpy()

                self.last_update = time.time()

                print(
                    f"Step {self.step:4d} / {self.N:4d} ({delta_step_best}) || "
                    f"beta: {self.beta:.1f}, "
                    f"mutations/s: {fps:.2f}, "
                    f"bad/good_accepts: {np.sum(self.bad_accepts[-100:])}/{np.sum(self.good_accepts[-100:])}",
                    flush=True,
                )
                print(f"STATS      || loss: {E_curr:.3f}, bkg: {background_loss:.3f}, mtf: {mtf_loss:.3f}, site: {site_loss:.3f}")
                print(f"BEST STATS || loss: {self.best_E:.3f}, bkg: {b_bkg:.3f}, mtf: {b_mtf:.3f}, site: {b_site:.3f}")

                #timings
                total_time = time.time() - start_time
                print(f"total time: {total_time:.1f}")
                print(f"mut time:   {muttime:.2f}\t| {100*muttime/total_time:.2f}%")
                print(f"nn time:    {nntime:.1f}\t| {100*nntime/total_time:.2f}%")
                print(f"loss time:  {losstime:.2f}\t| {100*losstime/total_time:.2f}%")
                print(f"score time: {scoretime:.1f}\t| {100*scoretime/total_time:.2f}%")
                print(f"misc time:  {misctime:.2f}\t| {100*misctime/total_time:.2f}%")

                plot_muts(self.gradient, self.results_dir / "gradient.jpg")
                plot_muts(self.good_mutations, self.results_dir / "mutations.jpg")

                if time.time() - self.last_graph > 1.8*cfg.report_interval or self.step == 0:
                    self.last_graph = time.time()

                    distogram = distogram_distribution_to_distogram(
                       self.best_distogram_distribution.cpu().detach().numpy()
                    )
                    plot_distogram(
                       distogram,
                       self.results_dir
                       / "distogram_evolution"
                       / f"{self.step:06d}_{self.best_E:.4f}.jpg",
                       clim=cfg.limits["dist"],
                    )

                    plot_progress(
                        E_tracker,
                        self.results_dir / "progress.jpg",
                        title=f"Optimization curve after {self.step} steps",
                    )
                    print(f"\n--- Current best: {self.best_sequence}")

            if (1 + self.step) % (cfg.report_interval//2) == 0:
                with open('control.txt', 'r') as reader:
                    lines = reader.readlines()
                    line = lines[0].strip()
                if line == "eval":
                    cmd = lines[1].strip()
                    print("EVALUATE")
                    print("cmd: " + cmd)
                    try:
                        eval(cmd.strip())
                    except:
                        print("ERROR")
                if line == "exit" or line == "break":
                    print("ending due to command")
                    break
                elif line == 'matrix':
                    if lines[1].strip() == 'n': cfg.MATRIX = False
                    else: cfg.MATRIX = True
                    print("MATRIX: " + str(cfg.MATRIX))
                elif line == 'gradient':
                    if lines[1].strip() == 'n': cfg.GRADIENT = False
                    else: cfg.GRADIENT = True
                elif line == "beta":
                    try:
                        value = float(lines[1].strip())
                        self.beta = value
                    except: print("input error")
                elif line == "bkw":
                    try:
                        value = float(lines[1].strip())
                        self.bkg_weight = value
                    except: print("input error")
                elif line == "aaw":
                    try:
                        value = float(lines[1].strip())
                        self.aa_weight = value
                    except: print("input error")
                elif line == "mw":
                    try:
                        value = float(lines[1].strip())
                        self.motif_weight = value
                    except: print("input error")
                elif line == "pause":
                    while True:
                        with open('control.txt', 'r') as reader:
                            line = reader.readlines()[0].strip()
                        if line != "pause":
                            break
                        else:
                            print("pause...")
                            time.sleep(30)

            if (delta_step_best > 300 and self.step > 1500 and self.step % 100 == 0) or (cfg.FAST and self.step > 499 and self.step % 50 == 0):
                std = np.std(np.array(E_tracker)[-1000:])
                n_std = std / np.array(E_tracker)[-1000:].mean()
                print("sd: " + str(std))
                if abs(std) < 0.005 or (cfg.FAST and abs(std) < 0.02):
                    break

            if self.step % self.M == 0 and self.step != 0:
                seq = self.fixup_MCMC(seq)

            misctime += time.time() - misc_start

        ########################################

        # Save final results before exiting:
        seq_curr = torch.from_numpy(
            aa2idx(self.best_sequence).reshape([1, self.seq_L])
        ).long()
        model_input, msa1hot = prep_seq(seq_curr)
        structure_predictions = self.structure_models(model_input)
        E_curr, self.best_metrics = self.loss(
            seq_curr, structure_predictions, msa1hot, track=True
        )

        #for key in self.best_metrics.keys():
        #    self.best_metrics[key] = v(metrics[key])

        TM_score_proxy = top_prob(structure_predictions['dist'], verbose=False)
        TM_score_proxy = TM_score_proxy[0]  # We're running with batch_size = 1

        self.best_metrics["TM_score_proxy"] = TM_score_proxy
        self.best_metrics["sequence"] = self.best_sequence
        self.best_metrics["motifweight"] = self.motif_weight
        self.best_metrics["motifmode"] = self.motifmode
        self.best_metrics["steps"] = self.step
        self.best_metrics["motifs"] = "\"" + str(self.motifs.copy()) + "\""

        # Dump distogram:
        best_distogram_distribution = structure_predictions['dist'].detach().cpu().numpy()
        distogram = distogram_distribution_to_distogram(best_distogram_distribution)
        plot_distogram(
            distogram,
            self.results_dir / f"result.jpg",
            clim=cfg.limits["dist"],
        )
        plot_progress(
            E_tracker,
            self.results_dir / "progress.jpg",
            title=f"Optimization curve after {self.step} steps",
        )

        # Write results to csv:
        with (self.results_dir / "result.csv").open("w") as f:
            for key, val in self.best_metrics.items():
                f.write(f"{key},{val}\n")

        with (self.results_dir / "mutation_log.csv").open("w") as f:
            for mut in self.mutation_log:
                f.write(f"{mut}\n")

        return self.best_metrics


#Main file for namsless should splinter later in development

#UNITS
#time:s
#weight:g
#length:cm
#density: g/cm^3

import numpy as np
import scipy.signal
import scipy.optimize
import matplotlib.pyplot as plt
from collections import OrderedDict
import os
import datetime
import sys

SMALL_SIZE = 16
MEDIUM_SIZE = 18
BIGGER_SIZE = 24

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

def pull_spectrum(spec_file):
    """
    Function to extract the spectrum from a *.Spe file
    Parameters:
        spec_file - file name of *.Spe file
    Returns:
        spectrum - numpy array of the number of counts per bin
    """
    with open(spec_file) as data:
        spectrum = np.array([float(i.strip()) for i in data.readlines()[12:2060]])
    return spectrum

def smooth(y, box_pts):
    """
    Simple 1-D convolution function
    Parameters:
        y - 1-D numpy array to be smoothed
        box_pts - width of kernel
    Returns:
        y_smooth - smoothed version of y
    """
    box = np.ones(box_pts)/box_pts
    y_smooth = np.convolve(y, box, mode='same')
    return y_smooth

class ATen:
    def __init__(self,input_filepath):
        self.banner = """                     _______     _________ _______  _
                    (  ___  )    \__   __/(  ____ \( (    /|
                    | (   ) |       ) (   | (    \/|  \  ( |
                    | (___) | _____ | |   | (__    |   \ | |
                    |  ___  |(_____)| |   |  __)   | (\ \) |
                    | (   ) |       | |   | (      | | \   |
                    | )   ( |       | |   | (____/\| )  \  |
                    |/     \|       )_(   (_______/|/    )_)

               "The People's 1D Continuous Energy Attenuation Code"

                             ...................
                             .. . ..............
             ......................,.,~$OOZ7O+.....................
             .....................Z$ZOO8OOZDMZ~....................
             .....................I8DDD8O88D888I...................
             ......................MNMNDDD888888:..................
             .......................$MMMMMNN888O8..................
             ........................=MNNNDDDDOOOI.................
             .........................=DDDDD88OZ88.................
             .........................ID88DD88OO88$................
             ......................,.$8DD8DDDD8OO88I...............
             ....................~OD88DN8DDD8DDDOO88...............
             ...............,IOOZ8888M=+DD88O888DOO8Z..   .........
             ............?O$OO8DNNMZ,..$DD8OO8888D8OZZZ   .........
             .........:ZO8DDMM8........DD888O8888OOZM.... .........
             ....,=8MNNNM...... .......ND888888D8ZD8D..............
             ... ...?I?7,,.............DD8D8888D8D88N..............
             ...   .... .    ..........+DDDDDDDDDD8Z8:.............
             ..........................DDDDDDDDDDDD8O8,............
             .........................88DDNDNDDDNNN8O88+...........
             ........................O8D8DDNNNNNNNNDO88D...........
             .......................+8DDDNNNMNMNNNNDOO887..........
             ......................,NDDNNNM.MMNNODNNOO88D..........
             .....................=8DNMN=:8NNDOO8NMMDOO8N:.........
             ....................:8DN~.......?:.....NO88D,.........
             ....................Z88ND.... ..+......,NO8D..........
             ................... .D888M. ....7.......Z888..........
             ......................OD8DM,....7......7OO8:..........
             ........................DDOM...$8$:....OOO8:..........
             ................... ......IODMMMDIM7..+OO88...........
             .........................~N8NNN8M7DMNMOOODN...........
             ........................D8M$N=DD:Z$NNM8OODO...........
             ...................  .,NDONDDD$8:ZM..M888N:...........
             ....................  D8ZM..8.$O,8I.ZDOO8D............
             ................... .?DDO7..?D$8+D.INDO888............
             ....................,DNO.MI~.IID==.=.8O8DN+. .........
             ................  ..O8$7..,D,.DM?~M.+OODMN8...........
             ................  ..DND.D$...N$$ON.Z:O8=8DM  .........
             ................... 8NN .:$DDIZDON,~8ZD$D8M ..........
             ....................8DI.  , =DND$$$,88I.O8M...........
             ....................DDO $8?..N8DNZ.8ZDNID8D ..........
             ....................8N8 ...M+?=DM$.OO?.8DD$ ..........
             ....................MNN,7ZD $N$+N8ZZO:.DDM.. .........
             ....................ODO8D..=D.M.8+DON.Z8O+ ...........
             ................... .MN$=.ZI..O,N?8D~~NOM.............
             ...................  :DDDZ~I.8$:D,D$8NDM..............
             ..................... 7ND8D.,7..ODDDNNMO..............
             ................... . ..MDN8NM8MMMDDD7................
             ........................ :MNDDDDMNMI .................
             ...................... ....  ~$,:.... ................
                                   ................

"""

        self.input_filepath = input_filepath
        self.casename = self.input_filepath.split("/")[-1][:-3]
        self.working_directory = "./" + "/".join(self.input_filepath.split("/")[:-1]) + "/"
        self.script_dir = "/".join(os.path.realpath(__file__).split("/")[:-1]) + "/"

        os.chdir(self.working_directory)

        #Reading contecnts of input file
        with open(input_filepath) as input_file:
            self.raw_input = input_file.readlines()
            input_data = [line.strip() for line in self.raw_input if not (line[0] == "#" or line == "\n")]

        #Read raw input text into cards
        self.material_card = input_data[input_data.index('READ MATERIAL')+1:input_data.index('END MATERIAL')]
        self.geometry_card = input_data[input_data.index('READ GEOMETRY')+1:input_data.index('END GEOMETRY')]
        self.parameters_card = input_data[input_data.index('READ PARAMETERS')+1:input_data.index('END PARAMETERS')]
        self.paths_card = input_data[input_data.index('READ PATHS')+1:input_data.index('END PATHS')]

        #initializing material card
        self.material_process()

        #parsing geometry related inputs
        self.layer_materials = [a for a in self.geometry_card if a.split()[0] == "layer_material"][0].split()[2:]
        self.layer_thicknesses = np.array([a for a in self.geometry_card if a.split()[0] == "layer_thicknesses"][0].split()[2:]).astype(np.float32)

        #parsing parameter related inputs
        self.parameters = {x[0] : x[1] for x in [parameter.replace(" ", "").split("=") for parameter in self.parameters_card]}

        #parsing path related inputs
        self.paths = {x[0] : self.working_directory + x[1] for x in [parameter.replace(" ", "").split("=") for parameter in self.paths_card]}

        self.bin_calibration(self.paths["cs137spec_filepath"])

        self.source_spec = pull_spectrum(self.paths["source_filepath"])/float(self.parameters["source_time"])

        if "background_filepath" in self.paths:
            self.source_spec -= pull_spectrum(self.paths["background_filepath"])/float(self.parameters["background_time"])

        #identigy groups and strengths and assign "cross sections"
        self.id_groups(plot_spec_peaks = True, plot_norm_peaks = True)
        self.ac_process()

    def bin_calibration(self,cs137spec_filepath):
        """
        Function to set energy values to each bin. Must be based on cs137 spectrum.
        Parameters:
            self - does not pull any attributes from self
            cs137spec_filepath - file name  of a *.Spe file containing unattenuated spectrum from a Cs-137 source
        Returns:
            self.bin_energies - numpy array containing the estimated energies associated with each bin
        """
        def how_much_gaussian(n, x, y):
            def gaus(x,a0,mu,sigma):
                return a0*(sigma*np.sqrt(2*np.pi))**-1*np.exp(-(x-mu)**2/(2*sigma**2))

            a0 = np.trapz(y)
            mean0 =  x[y.argmax()]
            sigma0 =  sum(y*(x-mean0)**2)/n

            popt,pcov = scipy.optimize.curve_fit(gaus,x,y,p0=[a0,mean0,sigma0],maxfev=1000000)
            perr = np.sqrt(np.diag(pcov)).sum()

            if perr > 500:
                return False
            else:
                return True

        cs137_peak = 661.6 #keV

        #Smooth the spectrum and extract energy local maxima
        cs137_spec = pull_spectrum(cs137spec_filepath)
        plt.show()

        cs137_spec = smooth(cs137_spec,30)
        peaks, empty_dict = scipy.signal.find_peaks(cs137_spec)

        #calculate the prominence associated with each maxima and find the bin with the largest prominence
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(cs137_spec, peaks)
        peaks = np.compress(prominences > prominences.mean(), peaks)
        left_bases = np.compress(prominences > prominences.mean(), left_bases)
        right_bases = np.compress(prominences > prominences.mean(), right_bases)

        gauss_errors = np.array([])
        for i in range(peaks.size):
            n = right_bases[i]-left_bases[i]
            x = np.linspace(left_bases[i],right_bases[i],n)
            y = cs137_spec[left_bases[i]:right_bases[i]]
            gauss_errors = np.append(gauss_errors, how_much_gaussian(n,x,y))

        peaks = np.compress(gauss_errors, peaks)
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(cs137_spec, peaks)

        #Find bin of maximum energy
        maxEbin = peaks[prominences.argmax()]
        self.maxEbin = maxEbin
        self.cs137_spec = cs137_spec
        delE = cs137_peak/maxEbin

        #Create bin energy array
        self.bin_energies = np.arange(0, cs137_spec.size * delE, delE)

        return 1

    def id_groups(self, plot_spec_peaks = False, plot_norm_peaks = False, width_threshold = [30, 170], width_rel_height = 0.7):
        """
        Function to identify energy group and count rate of each energy group.
        Parameters:
            self - both bin_energies and source_spec must be defined
            plot_spec_peaks - boolean to plot full source spectrum with e group identified
            plot_norm_peaks = boolean to plot norm fits as well. Should only be true if plot_spec_peaks is true
            width_threshold - tuning parameter for acceptable peak width
        Returns:
            self.group_counts - count rate associated with each group
            self.group_energies - energy value associated with each group
        """
        #Define gaussian checker for later in function
        def is_gaussian(n, x, y, width_rel_height,be,plot_norm_peaks):
            def gaus(x,a0,mu,sigma):
                return a0*(sigma*np.sqrt(2*np.pi))**-1*np.exp(-(x-mu)**2/(2*sigma**2))

            a0 = np.trapz(y) #4253
            mean0 =  x[y.argmax()] #sum(x*y)/n #377.977
            sigma0 =  sum(y*(x-mean0)**2)/n # 114.967

            popt,pcov = scipy.optimize.curve_fit(gaus,x,y,p0=[a0,mean0,sigma0],maxfev=100000)
            perr = np.sqrt(np.diag(pcov)).sum()

            if perr > 30 or perr is np.nan:
                return (False, 0)
            else:
                if plot_norm_peaks:
                    plt.plot(be,gaus(x,*popt),'o', c = "mediumseagreen", MarkerSize=5, label = "Gaussian Approximation")
                return (True, popt[0])

        #initializing figure if necessary
        if plot_spec_peaks or plot_norm_peaks:
            plt.figure(figsize=[15,8])

        #Smooth signal to and easily identify spectrum peaks
        smoothed_spec = self.source_spec
        peaks, empty_dict = scipy.signal.find_peaks(smoothed_spec)

        #Apply criteria that prominences must be greater than the average promenance
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(smoothed_spec, peaks)
        peaks = np.compress(prominences > prominences.mean(), peaks)

        #Apply criteria that widths must be greater than width_threshold
        widths, width_heights, leftips, rightips = scipy.signal.peak_widths(x = smoothed_spec, peaks = peaks, rel_height = width_rel_height)
        peaks = np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), peaks)
        leftips = np.floor(np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), leftips)).astype(np.int32)
        rightips = np.ceil(np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), rightips)).astype(np.int32)
        widths = np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), widths)

        #Apply criteria that shape must be appropriately gaussian
        #Hiding in here is getting the count rate under each peak as a return from the is_gaussian
        remove_i = np.array([]).astype(int)
        group_counts = np.array([])
        for i in range(peaks.size):
            n = rightips[i]-leftips[i]
            x = np.linspace(leftips[i],rightips[i],n)
            y = self.source_spec[leftips[i]:rightips[i]]
            gaussianality, group_count = is_gaussian(n,x,y, width_rel_height,self.bin_energies[leftips[i]:rightips[i]],plot_norm_peaks = plot_norm_peaks)
            if not gaussianality:
                remove_i = np.append(remove_i, i)
            else:
                group_counts = np.append(group_counts, group_count)

        peaks = np.delete(peaks,remove_i)
        leftips = np.delete(leftips,remove_i)
        rightips = np.delete(rightips,remove_i)

        #Add optional peak plotted to make sure peaks were done correctly
        if plot_spec_peaks:
            plt.plot(self.bin_energies,self.source_spec, c = "royalblue",label = "Source Spectrum")

            for i in range(leftips.size):
                plt.axvline(self.bin_energies[peaks[i]], c = "crimson", label = "Group Energy Value")

            plt.ylabel("Counts [#/s]")
            plt.xlabel("Energy [keV]")
            handles, labels = plt.gca().get_legend_handles_labels()
            by_label = OrderedDict(zip(labels, handles))
            plt.legend(by_label.values(), by_label.keys())
            plt.tight_layout()

        self.group_counts = group_counts
        self.group_energies = self.bin_energies[peaks]
        return 1

    def material_process(self):
        """
        Function to decompose entries in MATERIAL CARD from text into data structs
        Parameters:
            self - must have material_card specified
        Returns:
            self.materials - dictionary-containing two items, a density and a
            material composition list
        """
        self.materials = {}
        for entry in self.material_card:
            entry = entry.split()
            material = entry.pop(0)
            density = float(entry.pop(0))
            composition = [[int(entry[2*i]), float(entry[2*i-1])] for i in range(int(len(entry)/2))]
            self.materials[material] = [density, composition]

        return 1

    def ac_process(self):
        """
        Function to assign cross sections to each energy group for each material
        Parameters:
            self - must have self.materials defined
        Returns:
            self.ac - dictionary with keys of each material ID and values of list representing energies
                    - energy list as follows
                        - each row is a separate nuclide
                        - each column is an energy group
                        - sum each row for total ac for an energy group
        """
        def pull_ac(filename):
            #function to pull cross sections from libraries in workable form
            with open(filename, "r") as f:
                lib = np.array([s[3:].split() for s in f.readlines()])[:,:2].astype(np.float32)
            return lib

        self.ac = {}
        for key, value in self.materials.items():
            density = value[0]
            group_ac = np.empty((len(value[1])))
            for energy in self.group_energies:
                ac = np.array([])
                for znum, mass_frac in value[1]:
                    library = self.script_dir + "alib/" + [i for i in os.listdir(self.script_dir + "alib/") if "alib_" + str(znum) in i][0]
                    temp_ac = pull_ac(library)
                    ac = np.append(ac, np.interp(energy/1000,temp_ac[:,0],temp_ac[:,1])*mass_frac)

                group_ac = np.vstack((group_ac, ac*density))
            self.ac[key] = group_ac[1:,:] #some adjustment that needs to be made because of array initialization
        return 1

    def compute(self):
        def solve_master(master_ara, div_rows):
            for i in range(master_ara.shape[0]):
                if np.isclose(master_ara[i,1],0):
                    thickness = master_ara[i,2] - master_ara[i-1,2]
                    master_ara[i,1] = master_ara[i-1,1]*np.exp(-master_ara[i-1,3]*thickness)

        #setting up depth vector and n depending on "layer_mesh_divs"
        if "layer_mesh_divs" in self.parameters.keys():
            depth = np.array([0])
            material_list = []

            n = int(self.parameters["layer_mesh_divs"])

            for thickness, material in zip(self.layer_thicknesses,self.layer_materials):
                material_list += n*[material]
                start = depth[-1]
                stop = depth[-1]+thickness
                depth = np.append(depth, np.linspace(start + (stop-start)/(n+1) ,depth[-1]+thickness,n))
        else:
            n = 1
            depth = np.append(0,np.cumsum(self.layer_thicknesses))
            material_list = self.layer_materials

        #set up master array with columns energy, strength, depth, and atten coeff
        master_ara = np.zeros((depth.size*self.group_energies.size, 4))

        #create array with enery group repeated the same number of times as the length of depth vec
        working_grp = np.array([])
        for i in reversed(range(self.group_energies.size)):
            working_grp = np.append(working_grp, np.tile(self.group_energies[i],depth.size))

        #set up attenuation coeff vector
        ac_vec = np.array([])

        self.material_divisions = []

        for i in range(self.group_energies.size):
            ac_vec = np.append(0,ac_vec)
            for mat in material_list:
                ac_vec = np.append( self.ac[mat][i,:].sum(), ac_vec)
                self.material_divisions.append(mat)

        #place these vectors in master_ara
        master_ara[:,0] = working_grp
        master_ara[:,2] = np.tile(depth, self.group_energies.size)
        master_ara[:,3] = ac_vec

        #finding rows that start each energy group
        div_rows = np.array([0])
        for i in range(master_ara.shape[0]-1):
            if master_ara[i,3]==0:
                div_rows = np.append(div_rows, i+1)

        #initializing count rate for each energy group
        master_ara[div_rows,1] = self.group_counts

        #Fill in counts to each division in master_ara
        solve_master(master_ara, div_rows)

        self.master_ara = [master_ara,div_rows]

        return 1

    def print_output(self):
        def heading(string):
            return("\n" + 80*"-" + "\n" + string + "\n" + 80*"-" + "\n")

        #set up output and banner
        output = open(self.working_directory + self.casename + ".out", "w")
        output.write(self.banner + "\n" + str(datetime.datetime.now()) + "\n\n")

        #printing raw input
        output.write(heading("Raw Input"))
        output.write("".join(self.raw_input))

        #printing cards
        output.write(heading("Material Card"))
        output.write("\n".join(self.material_card) + "\n")
        output.write(heading("Geometry Card"))
        output.write("\n".join(self.geometry_card) + "\n")
        output.write(heading("Parameters Card"))
        output.write("\n".join(self.parameters_card) + "\n")
        output.write(heading("Paths Card"))
        output.write("\n".join(self.paths_card) + "\n")

        #printing materials
        output.write(heading("Materials"))
        for key, item in self.materials.items():
            #print(key,item[)
            output.write("Material: %s       Density: %.4f [g/cm^3]\n" % (key.ljust(15), item[0]))
            output.write("Atomic Number        Mass Fraction\n")
            for z,p in item[1]:
                output.write("%i                   %.5f\n"%(z,p))
            output.write("\n")

        #printing geomety
        output.write(heading("Geometry"))
        output.write("Material             Layer Thickness [cm]    Cumulative Thickness [cm]\n")
        accum = 0
        for i in range(len(self.layer_materials)):
            accum += self.layer_thicknesses[i]
            output.write("{:<17}".format(self.layer_materials[i]) + "    %.5f"%self.layer_thicknesses[i] + " "*17 + "%.5f"%(accum) + "\n")

        #printing group energy and strength
        output.write(heading("Energy Group Identification"))
        output.write("Group Energy [keV]                Group Count Rate [#/s]\n")
        for i in range(len(self.group_counts)):
            grp = "%.5f"%(self.group_energies[i])
            output.write(grp + (34 - len(grp))*" " + "%.5f\n"%(self.group_counts[i]))

        #printing master array
        output.write(heading("Master Array"))
        output.write("Grp Energy [keV]   Count Rate [#/s]   Depth [cm]   Atten Coeff [cm^2]   Material\n")

        i = 0
        for j, row in enumerate(self.master_ara[0]):
            e = "%.3f"%row[0]
            s = "%.3f"%row[1]
            d = "%.3f"%row[2]
            c = "%.5f"%row[3]

            if np.isclose(row[3],0):
                m = ""
            else:
                m = self.material_divisions[i]
                i += 1

            output.write(e + " "*(19-len(e)) + s + " "*(19-len(s)) + d + " "*(13-len(d))+ c + (21-len(c))*" " + m + "\n")

            if m == "":
                output.write(26*" - " + "\n")

        #attenuation coefficient libraries
        output.write(heading("Attenuation Coefficient Values"))
        for key, item in self.ac.items():
            output.write("Material = {mat}\n".format(mat = key))
            output.write("Atomic Number   Mass Fraction    ")
            for i, e in enumerate(self.group_energies):
                a = "E%i:%.2f"%(i,e)
                output.write(a + " "*(15-len(a)))
            output.write("\n")
            for i in range(item.T.shape[0]):
                anum = str(self.materials[key][1][i][0])
                frac = "%.4f"%self.materials[key][1][i][1]
                output.write(anum + " "*(16-len(anum)) + frac + " "*11 )
                for j in range(item.T.shape[1]):
                    const = "%.4f"%item.T[i,j]
                    output.write(const + " "*(15-len(const)))
                output.write("\n")
            output.write("\n")

        #printing calibration spectra with peak identified
        output.write(heading("Cs-137 Calibration Spectrum"))
        output.write("Bin Number           Count Number\n")
        for i in range(self.cs137_spec.size):
            output.write("{:<21}".format(str(i)) + str(self.cs137_spec[i]))

            if i == self.maxEbin:
                output.write("    *661.5 keV peak, used for calibration")

            output.write("\n")

        #printing source spectra
        output.write(heading("Unattenuated Spectra"))
        output.write("Bin Energy [keV]              Count Rate [#/s]\n")
        for i in range(self.source_spec.size):
            e = "%.5f"%self.bin_energies[i]
            s = "%.5f\n"%np.abs(self.source_spec[i])
            output.write(e +(30-len(e))*" " + s)

        output.close()

if __name__ == '__main__':
    ms_input = sys.argv[2]

    test = ATen(ms_input)
    test.compute()
    test.print_output()

    plt.show()


"""MCEq Flux Models

This script implements the use of MCEq flux models via the IceCube standard
method of getFlux(ptype, energy, costheta). As such it may be used as a
drop-in replacement for other fluxes. Weighting in IceCube is performed
by multiplying the flux by the normalized one weight:

NuGen:
    (with generator)
        weight = p_int * (flux_val / unit) * generator(energy, ptype, costheta)
    (without generator)
        weight = flux_val * one_weight / (type_weight * n_events * n_files)

    with flux_val = flux_object.getFlux(ptype, energy, costheta)

It is recommended to cache the results of MCEq because these take a while
to produce. By default, the cache file is chosen
to be located in the 'resources' directory relative to the location of this
script. You may also set the environment variable 'MCEQ_CACHE_DIR' in order
to choose a different location for the cache file, or pass in an explicit
cache file when initializing the MCEQFlux object.

Environment Variables:

    'MCEQ_CACHE_DIR':
        If provided, the MCEq cache file will be written to this directory.

    'MKL_PATH':
        Path to the MKL libraries. If provided, these are passed on to MCEq.
        Note: the python package can be installed directly with MKL support
        via 'pip install MCEq[MKL]'.

Credit for the vast majority of code in this file goes to Mathis Boerner.
"""
import os
import logging
from copy import deepcopy
import os
import numpy as np
from scipy.interpolate import RectBivariateSpline

import ic3_labels

log = logging.getLogger('MCEqFlux')

# If cashier is available, set up directory for caching of MCEq results
try:
    from ic3_labels.weights.resources.cashier import cache
    got_cashier = True

    if 'MCEQ_CACHE_DIR' in os.environ:
        cache_dir = os.environ['MCEQ_CACHE_DIR']
        log.info("Found 'MCEQ_CACHE_DIR' in environment variables: {}".format(
            cache_dir))

        if not os.path.exists(cache_dir):
            log.info('Creating cache directory: {}'.format(cache_dir))
            os.makedirs(cache_dir)

        CACHE_FILE = os.path.join(cache_dir, 'mceq.cache')

    else:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        CACHE_FILE = os.path.join(script_dir, 'resources', 'mceq.cache')

    log.info('Using MCEq cache file: {}'.format(CACHE_FILE))

except ImportError:
    got_cashier = False
    CACHE_FILE = None
    log.info("Could not import 'cashier'. MCEq results will not be cached!")

# Dictionary that converts ptype -> MCEq type string
PTYPE_CONVERTER = {
    12: 'nue',
    -12: 'antinue',
    14: 'numu',
    -14: 'antinumu',
    16: 'nutau',
    -16: 'antinutau',
}

def get_spline(
        interaction_model,
        primary_model,
        months,
        theta_grid,
        theta_grid_cos,
        cached=True,
        cache_file=CACHE_FILE,
        cache_read_only=False):
    """Get MCEq spline

    Solves the MCEq cascade equations for the given parameters. The equations
    are solved on the provided grid and interpolated.

    Parameters
    ----------
    interaction_model : str
        The interaction model. This is passed on to `MCEqRun`.
    primary_model : str
        The primary model to use. Must be one of:
            GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
    months : list of str
        The months for which to solve the cascade equations. These must be
        provided as a list of month names, e.g. ['January', 'August']. A list
        of splines will be returned of the same length as `months`.
    theta_grid : array_like
        The grid points in theta to evaluate on in degrees. If `theta_grid_cos`
        is True, this is instead cos(theta).
    theta_grid_cos : bool
        If True, `theta_grid` is interpreted as cos(theta), i.e. arccos() is
        applied first.
    cached : bool, optional
        If True, the result will be cached, or taken from cache if previously
        already computed. This is recommended, as MCEq takes a while to run.
    cache_file : str, optional
        The path to the cache file to use.
    cache_read_only : bool, optional
        If True, the cache is read only.

    Returns
    -------
    dict
        The result of MCEq together with the fitted splines. The structure is
        as follows:
        {
            # first month provided via `months`
            0: {
                'total_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'total' flux. The dictionary keys are the
                    PDG particle encodings.
                'conv_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'conv' flux. The dictionary keys are the
                    PDG particle encodings.
                'pr_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'pr' flux. The dictionary keys are the
                    PDG particle encodings.
                'total_flux_dict': dict of array_like
                    A dictionary with the total flux for each grid point.
                    This is the result obtained from MCEq for the 'total' flux.
                'conv_flux_dict': dict of array_like
                    A dictionary with the conv flux for each grid point.
                    This is the result obtained from MCEq for the 'conv' flux.
                'pr_flux_dict': dict of array_like
                    A dictionary with the prompt flux for each grid point.
                    This is the result obtained from MCEq for the 'pr' flux.
                'config_updates':   dict
                    A dictionary of config updates that were applied to
                    mceq_config prior to solving the equations.
                'mceq_version' : str
                    The MCEq version that was used to create the splines.
                'ic3_labels_version' : str
                    The version of the ic3-labels package that was used to
                    create the splines.
                'e_grid' : array_like
                    The grid of energy points in log10.
                'theta_grid' : array_like
                    The grid of thetas.
            }

            # second month provided via `months`
            1: {
                ...
            }

            ...
        }
    """
    log.info('Getting Spline for {}; {} (cached={})'.format(
        interaction_model,
        primary_model,
        cached))

    def __solve_month__(
            mceq_run, e_grid, theta_grid, theta_grid_cos,
            ptype_converter=PTYPE_CONVERTER,
            eps=1e-128,
            ):
        """Solve MCEq equations for the provided mceq_run instance.

        Parameters
        ----------
        mceq_run : MCEqRun instance
            The MCEqRun instance. This instance must be configured to use
            the desired geometry and season.
        e_grid : array_like
            The grid of energy points in log10.
        theta_grid : array_like
            The grid points in theta to evaluate on in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `theta_grid` is interpreted as cos(theta),
            i.e. arccos() is applied first.
        ptype_converter : dict, optional
            A dictionary that converts PDG encoding to MCEq type string.
        eps : float, optional
            A small float value > 0 that is used to clip the total flux
            prior to applying log10 for the spline fitting.

        Returns
        -------
        list of dict of RectBivariateSpline
            A list of dictionaries with the fitted splines for each particle
            type. The dictionary keys are the PDG particle encodings.
            The order of the dictionaries are: 'total', 'conv', 'pr'
        list of dict of array_like
            A list of dictionaries with the total flux for each grid point.
            This is the result obtained from MCEq.
            The order of the dictionaries are: 'total', 'conv', 'pr'
        """
        total_flux_dict = {}
        conv_flux_dict = {}
        pr_flux_dict = {}
        total_spline_dict = {}
        conv_spline_dict = {}
        pr_spline_dict = {}
        for key, value in ptype_converter.items():
            total_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))
            conv_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))
            pr_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))

        for i, theta_i in enumerate(theta_grid):
            if theta_grid_cos:
                theta_i = np.rad2deg(np.arccos(theta_i))
            mceq_run.set_theta_deg(theta_i)
            mceq_run.solve()

            # fill in flux totals
            for key, value in ptype_converter.items():
                total_flux_dict[key][:, i] = mceq_run.get_solution(
                    'total_{}'.format(value))
                conv_flux_dict[key][:, i] = mceq_run.get_solution(
                    'conv_{}'.format(value))
                pr_flux_dict[key][:, i] = mceq_run.get_solution(
                    'pr_{}'.format(value))

        # create splines
        for key, value in ptype_converter.items():
            total_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(total_flux_dict[key], eps, float('inf'))),
                s=0,
            )
            conv_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(conv_flux_dict[key], eps, float('inf'))),
                s=0,
            )
            pr_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(pr_flux_dict[key], eps, float('inf'))),
                s=0,
            )

        spline_dicts = [
            total_spline_dict, conv_spline_dict, pr_spline_dict
        ]

        flux_dicts = [
            total_flux_dict, conv_flux_dict, pr_flux_dict
        ]

        return spline_dicts, flux_dicts

    def __get_spline__(
            interaction_model,
            primary_model,
            months,
            theta_grid,
            theta_grid_cos,
            ):
        """Get MCEq spline for the provided settings

        Parameters
        ----------
        interaction_model : str
            The interaction model. This is passed on to `MCEqRun`.
        primary_model : str
            The primary model to use. Must be one of:
                GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
        months : list of str
            The months for which to solve the cascade equations. These must be
            provided as a list of month names, e.g. ['January', 'August']. A
            list of splines will be returned of the same length as `months`.
        theta_grid : array_like
            The grid points in theta to evaluate on in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `theta_grid` is interpreted as cos(theta),
            i.e. arccos() is applied first.

        Returns
        -------
        dict
            The result of MCEq together with the fitted splines.
            See documentation of `get_spline()` for more details.

        Raises
        ------
        AttributeError
            If the provided `primary_model` is unknown.
        """
        log.info('\tCalculating \'{}\' \'{}\''.format(
            interaction_model, primary_model))

        import mceq_config
        from MCEq import version
        from MCEq.core import MCEqRun
        import crflux.models as pm

        config_updates = {
            'h_obs': 1000.,
            'debug_level': 1,
        }
        if 'MKL_PATH' in os.environ:
            config_updates['MKL_path'] = os.environ['MKL_PATH']

        splines = {}
        pmodels = {
            "GST_3-gen": (pm.GaisserStanevTilav, "3-gen"),
            "GST_4-gen": (pm.GaisserStanevTilav, "4-gen"),
            "H3a": (pm.HillasGaisser2012, "H3a"),
            "H4a": (pm.HillasGaisser2012, "H4a"),
            "poly-gonato": (pm.PolyGonato, False),
            "TIG": (pm.Thunman, None),
            "ZS": (pm.ZatsepinSokolskaya, 'default'),
            "ZSP": (pm.ZatsepinSokolskaya, 'pamela'),
            "GH": (pm.GaisserHonda, None),
        }

        for i, month in enumerate(months):
            config_updates['density_model'] = (
                'MSIS00_IC', ('SouthPole', month))

            # update settings in mceq_config
            # Previous method mceq_config.config is deprecated and resulted
            # in pickle errors for deepcopy.
            for name, value in config_updates.items():
                setattr(mceq_config, name, value)

            try:
                pmodel = pmodels[primary_model]
            except KeyError:
                raise AttributeError(
                    'primary_model {} unknown. options: {}'.format(
                        primary_model, pmodels.keys()))
            mceq_run = MCEqRun(
                interaction_model=interaction_model,
                primary_model=pmodel,
                theta_deg=0.0,
                **config_updates)
            e_grid = np.log10(deepcopy(mceq_run.e_grid))
            spline_dicts, flux_dicts = __solve_month__(
                mceq_run,
                e_grid,
                theta_grid,
                theta_grid_cos)

            splines[i] = {}
            splines[i]['total_spline_dict'] = spline_dicts[0]
            splines[i]['conv_spline_dict'] = spline_dicts[1]
            splines[i]['pr_spline_dict'] = spline_dicts[2]
            splines[i]['total_flux_dict'] = flux_dicts[0]
            splines[i]['conv_flux_dict'] = flux_dicts[1]
            splines[i]['pr_flux_dict'] = flux_dicts[2]
            splines[i]['config_updates'] = deepcopy(config_updates)
            splines[i]['mceq_version'] = version.__version__
            splines[i]['ic3_labels_version'] = ic3_labels.__version__
            splines[i]['e_grid'] = e_grid
            splines[i]['theta_grid'] = theta_grid
        return splines

    if got_cashier and cached:
        if cache_file is None:
            cache_f = 'mceq.cache'
        else:
            cache_f = cache_file
        log.info('\tUsing cache \'{}\''.format(cache_f))

        @cache(cache_file=cache_f, read_only=cache_read_only)
        def wrapped_get_spline(
                interaction_model,
                primary_model,
                months,
                theta_grid,
                theta_grid_cos,
                ):
            return __get_spline__(
                interaction_model=interaction_model,
                primary_model=primary_model,
                months=months,
                theta_grid=theta_grid,
                theta_grid_cos=theta_grid_cos,
            )

        return wrapped_get_spline(
            interaction_model,
            primary_model,
            months,
            theta_grid,
            theta_grid_cos,
        )
    else:
        return __get_spline__(
            interaction_model=interaction_model,
            primary_model=primary_model,
            months=months,
            theta_grid=theta_grid,
            theta_grid_cos=theta_grid_cos,
        )

class MCEQFlux(object):

    """MCQe Flux Wrapper

    Attributes
    ----------
    min_theta_deg : float, optional
        The minimum value of the theta grid in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    max_theta_deg : float, optional
        The maximum value of the theta grid in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    month_weights : array_like
        A list of probabilities for each month. These are used as weights
        to sample the corresponding month for each MC events.
        These weights can, for instance, be set to the relative livetime in
        each month of the year. This will then account for seasonal variations.
    months : list of str
        A list of months for which the interpolation splines are created.
    random_state : np.random.RandomState
        The random state that is used to draw the month for each MC event.
    splines : dict
        A dictionary containing the MCEq result and fitted splines.
        See documentation of `get_splines()` for more details.
    theta_grid : array_like
        The grid points in theta to evaluate on in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    theta_grid_cos : bool
        If True, `min_theta_deg` and `max_theta_deg` are interpreted as
        cos(theta), i.e. arccos() is applied first.
    """

    def __init__(
            self,
            min_theta_deg=0.,
            max_theta_deg=180.,
            theta_grid_cos=False,
            theta_steps=181,
            season='full_year',
            flux_type='total',
            random_state=None,
            **kwargs):
        """Initialize MCEQFlux Instance

        Parameters
        ----------
        min_theta_deg : float, optional
            The minimum value of the theta grid in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        max_theta_deg : float, optional
            The maximum value of the theta grid in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `min_theta_deg` and `max_theta_deg` are interpreted as
            cos(theta), i.e. arccos() is applied first.
        theta_steps : int, optional
            The number of grid points between the specified min and max values.
        season : str, optional
            What season to use. This may either be a single month ,
            for example 'January', or 'full_year' may be used to run MCEq
            for every month of the year.
        flux_type : str, optional
            The flux type to compute. This must be one of
                'total': combined prompt and conv flux
                'pr': prompt neutrino flux
                'conv': conventional neutrino flux
            This will set the default flux type when calling `getFlux()`.
            You may, however, overwrite this defaul by passing an alternative
            flux type to `getFlux()`. Setting this default value allows for
            drop in replacement of other flux implementations in IceCube.
        random_state : np.random.Randomstate or int, optional
            An int or random state to set the seed.
        **kwargs
            Additional keyword arguments. (Not used!)
        """
        if not isinstance(random_state, np.random.RandomState):
            random_state = np.random.RandomState(random_state)
        self.random_state = random_state
        if season.lower() == 'full_year':
            self.months = [
                'January',
                'February',
                'March',
                'April',
                'May',
                'June',
                'July',
                'August',
                'September',
                'October',
                'November',
                'December',
            ]
        else:
            self.months = [season]

        if flux_type.lower() not in ['total', 'conv', 'pr']:
            raise ValueError('Flux type: {} must be on of {}'.format(
                flux_type.lower(), ['total', 'conv', 'pr']))

        self.flux_type = flux_type
        self.set_month_weights(np.ones_like(self.months, dtype=float))
        self.min_theta = min_theta_deg
        self.max_theta = max_theta_deg
        if theta_grid_cos:
            self.min_theta = np.cos(np.deg2rad(min_theta_deg))
            self.max_theta = np.cos(np.deg2rad(max_theta_deg))
        self.theta_grid_cos = theta_grid_cos
        self.theta_grid = np.linspace(
            self.min_theta, self.max_theta, theta_steps)

        self.splines = None

    def initialize(
            self,
            interaction_model='SIBYLL2.3c',
            primary_model='H3a',
            cached=True,
            cache_file=CACHE_FILE,
            cache_read_only=False,
            ):
        """Initialize MCEQFlux instance

        This will compute the splines or retrieve these from the cache
        if `cached` is True and if these have been previously computed.
        This method must be called prior to calls to `getFlux()`.

        Parameters
        ----------
        interaction_model : str
            The interaction model. This is passed on to `MCEqRun`.
        primary_model : str
            The primary model to use. Must be one of:
                GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
        cached : bool, optional
            If True, the result will be cached and if already computed, it will
            be retrieved from cache. This avoids recomputation of MCEq, which
            is recommended in order to reduce computation time.
        cache_file : str, optional
            The path to the cache file to use.
        cache_read_only : bool, optional
            If True, the cache is read only.
        """
        if cache_file is None:
            cache_file = CACHE_FILE

        self.splines = get_spline(
            interaction_model,
            primary_model,
            self.months,
            self.theta_grid,
            self.theta_grid_cos,
            cached=cached,
            cache_file=cache_file,
            cache_read_only=cache_read_only,
        )

        from MCEq import version

        # throw warning if there is a version mis-match.
        for key, spline in self.splines.items():
            msg = (
                'Cached file was created with {} version {}, '
                'but this is version {}!'
            )
            if version.__version__ != spline['mceq_version']:
                log.warning(msg.format(
                    'MCEq', spline['mceq_version'], version.__version__))

            if ic3_labels.__version__ != spline['ic3_labels_version']:
                log.warning(msg.format(
                    'ic3_labels',
                    spline['ic3_labels_version'],
                    ic3_labels.__version__,
                ))

    def set_month_weights(self, month_weights):
        """Summary

        Parameters
        ----------
        month_weights : array_like
            A list of probabilities for each month (these will be normalized
            internally). These are used as weights to sample the corresponding
            month for each MC events. These weights can, for instance, be set
            to the relative livetime in each month of the year.
            This will then account for seasonal variations.

        Raises
        ------
        AttributeError
            If the length of the provided `muon_weights` does not match the
            length of the specified months.
        """
        if len(month_weights) != len(self.months):
            raise AttributeError(
                'month_weights needs to be of the same '
                'length like self.months.'
            )
        self.month_weights = month_weights / np.sum(month_weights)

    def getFlux(
            self,
            ptype,
            energy,
            costheta,
            selected_month=None,
            random_state=None,
            flux_type=None,
            ):
        """Get flux for provided particle

        The flux is given in GeV^-1 cm^-2 s^-1 sr^-1 and may be used to
        weight NuGen events via the normalized `one_weight`:
            weight = flux * one_weight / (type_weight * n_events * n_files)

        Parameters
        ----------
        ptype : array_like or int
            The PDG encoding.
            For instance: I3MCWeightDict -> PrimaryNeutrinoType
        energy : array_like or float
            The energy of the primary particle.
            For instance: I3MCWeightDict -> PrimaryNeutrinoEnergy
        costheta : array_like or float
            The cos(zenith) angle of the primary particle.
            For instance: cos(I3MCWeightDict -> PrimaryNeutrinoZenith)
        selected_month : array_like, optional
            The month in which each event occurred. This must be given as
            an array of integer values between [0, 11] if `season` is
            'full_year'. If the `MCEQFlux` instance is initialized with only
            one month as the season, then `selected_month` must not be set.
            If None provided, the corresponding month of each event will be
            sampled via the defined `month_weights`.
        random_state : np.random.Randomstate or int, optional
            An int or random state to set the seed.
            If None provided, the random state will be used that was
            reated during initialization.
        flux_type : str, optional
            The flux type to compute. This must be one of
                'total': combined prompt and conv flux
                'pr': prompt neutrino flux
                'conv': conventional neutrino flux
            If None is provided, the specified default value at
            object instantiation time (__init__()) will be used.

        Returns
        -------
        array_like
            The flux for the given particle in GeV^-1 cm^-2 s^-1 sr^-1.

        Raises
        ------
        RuntimeError
            If MCEQFlux has not been initialized yet.
        ValueError
            If wrong `flux_type` is provided.
        """
        if self.splines is None:
            raise RuntimeError(
                'No splines calculated! Run \'initialize\' first')

        if len(self.months) == 1 and selected_month is not None:
            raise ValueError(
                'The months may not be set, since the MCEQFlux instance is '
                + 'initialized with only one month: {}'.format(self.months)
            )

        if flux_type is None:
            flux_type = self.flux_type
        elif flux_type.lower() not in ['total', 'conv', 'pr']:
            raise ValueError('Flux type: "{}" must be on of {}'.format(
                flux_type.lower(), ['total', 'conv', 'pr']))
            flux_type = flux_type.lower()

        if random_state is None:
            random_state = self.random_state
        elif not isinstance(random_state, np.random.RandomState):
            random_state = np.random.RandomState(random_state)

        # convert to numpy arrays and make sure these are at least 1D
        ptype = np.atleast_1d(ptype)
        energy = np.atleast_1d(energy)
        costheta = np.atleast_1d(costheta)

        if len(self.splines) > 1:
            if selected_month is None:
                int_months = np.arange(len(self.splines), dtype=int)
                selected_month = random_state.choice(
                    int_months,
                    replace=True,
                    size=len(energy),
                    p=self.month_weights,
                )
            else:
                selected_month = np.asarray(selected_month, dtype=int)
        else:
            selected_month = list(self.splines.keys())[0]
        flux = np.ones_like(energy)
        flux[:] = float('NaN')
        log10_energy = np.log10(energy)
        theta = np.rad2deg(np.arccos(costheta))

        for ptype_i in np.unique(ptype):
            mask_ptype = ptype == ptype_i
            for i in self.splines.keys():
                if isinstance(selected_month, int):
                    idx_ptype = mask_ptype
                else:
                    is_in_month = selected_month == i
                    idx_ptype = np.logical_and(mask_ptype, is_in_month)
                flux[idx_ptype] = self.splines[i][
                        flux_type + '_spline_dict'][ptype_i](
                    log10_energy[idx_ptype],
                    theta[idx_ptype],
                    grid=False)
        return np.power(10., flux)