#!/usr/bin/env python
# coding: utf-8

# ### Hill Climber algorithm to generate a team
#
# - takes in the list of available players for the gameweek with their predicted scores from the player model
# - constructs a 15-man squad adhering to budget connstraints and team constraints set by FPL

# In[ ]:

import numpy as np
import pandas as pd
import random as rd
import random
from random import randint
from random import sample
import matplotlib.pyplot as plt
import requests
pd.options.display.max_rows = None
pd.options.display.max_columns = None
import itertools

# In[78]:

class generate_team():

    def __init__(self, players):
        self.players = players
        self.initiate()

    def init_player_lists(self,players):
        #Method to initialise four vectors of 0, where each defines the position and is the length of the amount of
        #players in that position. Randomly assign 1's to an index in each list to initiate a squad
        #@params players = the list of available players
        #@return goalkeeper, defender, midfielder, attacker lists

        #make the lists from the actual players
        goalks, defs, mids, atts = [], [], [], []
        for i in players:
            if i[1] == 1:
                goalks.append(i)
            elif i[1] == 2:
                defs.append(i)
            elif i[1] == 3:
                mids.append(i)
            elif i[1] == 4:
                atts.append(i)

        #initialise each vector to have a random player, goalkeepers 2 players, def 5, mid 5, atk 3.
        goalkeepers = [0 for i in range(len(goalks))]
        x = [i for i in range(len(goalks))]
        random_numbers = random.sample(x, 2)
        for random_number in random_numbers:
            goalkeepers[random_number] = 1

        defenders = [0 for i in range(len(defs))]
        x = [i for i in range(len(defs))]
        random_numbers = random.sample(x, 5)
        for random_number in random_numbers:
            defenders[random_number] = 1

        midfielders = [0 for i in range(len(mids))]
        x = [i for i in range(len(mids))]
        random_numbers = random.sample(x, 5)
        for random_number in random_numbers:
            midfielders[random_number] = 1

        attackers = [0 for i in range(len(atts))]
        x = [i for i in range(len(atts))]
        random_numbers = random.sample(x, 3)
        for random_number in random_numbers:
            attackers[random_number] = 1

        return goalkeepers, defenders, midfielders ,attackers

    def init_hc(self,max_threshold, players):
        #Method to call initialise squad vectors. If the constraints are not met, repeat the initialisation
        #@params max_threshold = the budget constraint, players = the available players
        #@return goalkeeper, defender, midfielder, attacker lists

        goalkeeper, defender, midfielder, attacker = [], [], [], []
        for i in players:
            if i[1] == 1:
                goalkeeper.append(i)
            elif i[1] == 2:
                defender.append(i)
            elif i[1] == 3:
                midfielder.append(i)
            elif i[1] == 4:
                attacker.append(i)

        total_weight = max_threshold +1
        total_points = 51
        #total_weight cannot go over budget constraint
        while (total_weight > 100 or total_points < 50):
            gkp, defs, mids, atks = self.init_player_lists(players)
            g_weight = (sum([goalkeeper[i][4] for i in range(len(goalkeeper)) if gkp[i] == 1]))
            d_weight = (sum([defender[i][4] for i in range(len(defender)) if defs[i] == 1]))
            m_weight = (sum([midfielder[i][4] for i in range(len(midfielder)) if mids[i] == 1]))
            a_weight = (sum([attacker[i][4] for i in range(len(attacker)) if atks[i] == 1]))
            total_weight = g_weight + d_weight + m_weight + a_weight

            team_points = []
            team_points.append([goalkeeper[i][3] for i in range(len(goalkeeper)) if gkp[i] == 1])
            team_points.append([defender[i][3] for i in range(len(defender)) if defs[i] == 1])
            team_points.append([midfielder[i][3] for i in range(len(midfielder)) if mids[i] == 1])
            team_points.append([attacker[i][3] for i in range(len(attacker)) if atks[i] == 1])
            team_points = list(itertools.chain(*team_points))
            total_points = sum(team_points)
        return gkp, defs, mids, atks

    def evaluate_fitness(self,max_threshold, players, gkp, defs, mids, atks):
        #Method to evaluate fitness by looking at the weight of the squad. Similarly, to see if any team has
        #more than 3 players in the squad. This is not allowed in FPL rules
        #@params max-threshold = budget, players = available players, gkp,defs,mids,atks = squad player vectors
        #@return fitness = how many projected points, team_weight = cost of squad, team_points = indiv player points
        team_weight= []
        team_points = []
        fitness = 0

        goalkeeper, defender, midfielder, attacker = [], [], [], []
        for i in players:
            if i[1] == 1:
                goalkeeper.append(i)
            elif i[1] == 2:
                defender.append(i)
            elif i[1] == 3:
                midfielder.append(i)
            elif i[1] == 4:
                attacker.append(i)

        #calculate weight of squad (total cost)
        team_weight.append([goalkeeper[i][4] for i in range(len(goalkeeper)) if gkp[i] == 1])
        team_weight.append([defender[i][4] for i in range(len(defender)) if defs[i] == 1])
        team_weight.append([midfielder[i][4] for i in range(len(midfielder)) if mids[i] == 1])
        team_weight.append([attacker[i][4] for i in range(len(attacker)) if atks[i] == 1])
        team_weight = list(itertools.chain(*team_weight))

        #calculate total team points predicted
        team_points.append([goalkeeper[i][3] for i in range(len(goalkeeper)) if gkp[i] == 1])
        team_points.append([defender[i][3] for i in range(len(defender)) if defs[i] == 1])
        team_points.append([midfielder[i][3] for i in range(len(midfielder)) if mids[i] == 1])
        team_points.append([attacker[i][3] for i in range(len(attacker)) if atks[i] == 1])
        team_points = list(itertools.chain(*team_points))

        #calculate how many players are from each team to see if gone over
        team_code = []
        team_code.append([goalkeeper[i][2] for i in range(len(goalkeeper)) if gkp[i] == 1])
        team_code.append([defender[i][2] for i in range(len(defender)) if defs[i] == 1])
        team_code.append([midfielder[i][2] for i in range(len(midfielder)) if mids[i] == 1])
        team_code.append([attacker[i][2] for i in range(len(attacker)) if atks[i] == 1])
        team_code = list(itertools.chain(*team_code))
        overkill_team = False
        for team in range(1,21):
            i = 0
            for team_ in team_code:
                if team == team_:
                    i+= 1
                    if i > 3:
                        #overkill = too many players from one team
                        overkill_team = True

        #If team weight is too high or overkill is true, fitness of the squad is 0
        if sum(team_weight) > max_threshold or overkill_team is True:
            fitness = 0
        else:
            fitness = sum(team_points)
        return fitness, team_weight, team_points

    def mutate(self,mutation_rate, gkp, defs, mids, atks):
        #Method to mutate the squad in a certain way to find new better squads
        #@params mutation_rate = the likelihood of change, gkp,def,mids,atks = squad vectors
        #@return gkp,defs,mids,atks = new mutated squad

        a = rd.random()
        b = rd.random()
        c = rd.random()
        d = rd.random()

        #For each position, get a random number, swap out a current player and change a position in the vector
        #from zero to one to add a new player

        if a <= mutation_rate:
            x = [i for i in range(len(gkp))]
            random_number = random.sample(x, 1)
            present_numbers = [i for i in range(len(gkp))if gkp[i] == 1]
            #if player is already in the squad, dont do anything, if not carry on
            if random_number[0] not in present_numbers:
                #choose which player to remove randomly out of 2 players
                r = [i for i in range(0,2)]
                #remove the player
                remove_index = random.sample(r,1)
                for player in range(len(gkp)):
                    if player == present_numbers[remove_index[0]]:
                        gkp[player] = 0
                    if player == random_number[0]:
                        gkp[player] = 1

        if b <= mutation_rate:
            x = [i for i in range(len(defs))]
            random_number = random.sample(x, 1)
            present_numbers = [i for i in range(len(defs))if defs[i] == 1]
            #if player is already in the squad, dont do anything, if not carry on
            if random_number[0] not in present_numbers:
                #choose which player to remove randomly out of 5 players
                r = [i for i in range(0,5)]
                #remove the player
                remove_index = random.sample(r,1)
                for player in range(len(defs)):
                    if player == present_numbers[remove_index[0]]:
                        defs[player] = 0
                    if player == random_number[0]:
                        defs[player] = 1

        if c <= mutation_rate:
            x = [i for i in range(len(mids))]
            random_number = random.sample(x, 1)
            present_numbers = [i for i in range(len(mids))if mids[i] == 1]
            #if player is already in the squad, dont do anything, if not carry on
            if random_number[0] not in present_numbers:
                r = [i for i in range(0,5)]
                #choose which player to remove randomly out of 5 players
                remove_index = random.sample(r,1)
                #remove the player
                for player in range(len(mids)):
                    if player == present_numbers[remove_index[0]]:
                        mids[player] = 0
                    if player == random_number[0]:
                        mids[player] =1

        if d <= mutation_rate:
            x = [i for i in range(len(atks))]
            random_number = random.sample(x, 1)
            present_numbers = [i for i in range(len(atks))if atks[i] == 1]
            #if player is already in the squad, dont do anything, if not carry on
            if random_number[0] not in present_numbers:
                #choose which player to remove randomly out of 3 players
                r = [i for i in range(0,3)]
                #remove the player
                remove_index = random.sample(r,1)
                for player in range(len(atks)):
                    if player == present_numbers[remove_index[0]]:
                        atks[player] = 0
                    if player == random_number[0]:
                        atks[player] = 1

        return gkp, defs, mids, atks

    def hillclimber(self,gkp, defs, mids, atks, players, max_threshold, generations, mutation_rate):
        #Method to call hill climber algorithm with each individual method, used to rretrieve new better squads
        #@params gkp,defs,mids,atks = player vectors, max_threshold = budget, generations, mutation_rate
        #@return gkp,defs,mids,atks = player vectors, fitness = list of squad projected points over time

        fitness = []
        for i in range(generations):
            #if (i % 50000) == 0:
                #print(i/50000, "% completed")

            #make a copy of the exisiting squad
            gkp1 = gkp.copy()
            defs1 = defs.copy()
            mids1 = mids.copy()
            atks1 = atks.copy()

            #mutate copied squad, calc fitness of both new and old squads
            gkp1, defs1, mids1, atks1 = self.mutate(mutation_rate, gkp1, defs1, mids1, atks1)
            g0_fitness, team_weight, team_points = self.evaluate_fitness(max_threshold, players, gkp, defs, mids, atks)
            g1_fitness, team_weight, team_points = self.evaluate_fitness(max_threshold, players, gkp1, defs1, mids1, atks1)

            #if new squad has better fitness, this becomes the main squad
            if g1_fitness > g0_fitness:
                gkp = gkp1
                mids = mids1
                defs = defs1
                atks = atks1

            fitness.append(g0_fitness)

        return gkp, defs, mids, atks, fitness

    def graph(self,fitness,mutation_rate, players, gkp, defs, mids, atks,generations):
        #Method to call a visual representation of the graph
        #@params fitness = squad fitness, mutation_rate, players = available players, squad vectors, generations
        #@return total team points and total team weight

        #print("Mutation rate = ", mutation_rate)
        team_weight= []
        team_points = []

        goalkeeper, defender, midfielder, attacker = [], [], [], []
        for i in players:
            if i[1] == 1:
                goalkeeper.append(i)
            elif i[1] == 2:
                defender.append(i)
            elif i[1] == 3:
                midfielder.append(i)
            elif i[1] == 4:
                attacker.append(i)

        team_weight.append([goalkeeper[i][4] for i in range(len(goalkeeper)) if gkp[i] == 1])
        team_weight.append([defender[i][4] for i in range(len(defender)) if defs[i] == 1])
        team_weight.append([midfielder[i][4] for i in range(len(midfielder)) if mids[i] == 1])
        team_weight.append([attacker[i][4] for i in range(len(attacker)) if atks[i] == 1])
        team_weight = list(itertools.chain(*team_weight))

        team_points.append([goalkeeper[i][3] for i in range(len(goalkeeper)) if gkp[i] == 1])
        team_points.append([defender[i][3] for i in range(len(defender)) if defs[i] == 1])
        team_points.append([midfielder[i][3] for i in range(len(midfielder)) if mids[i] == 1])
        team_points.append([attacker[i][3] for i in range(len(attacker)) if atks[i] == 1])
        team_points = list(itertools.chain(*team_points))

        #print("Total final points:", sum(team_points))
        #print("Total final cost:", sum(team_weight))

        #plt.plot(range(generations), fitness, label = 'Fitness')
        #plt.legend()
        #plt.title('Fitness level of hill Climber over 500 generations')
        #plt.xlabel('Generations')
        #plt.ylabel('Fitness')
        #plt.show()

        return (sum(team_points)), (sum(team_weight))

    def project_squad(self,players, top_choice):
        #Method to get the best squad vector and change it from the labels back into the real projected points
        #@params players = available players, top_choice = top squad vectors
        #@return None

        goalks, defs, mids, atts = [], [], [], []
        for i in players:
            if i[1] == 1:
                goalks.append(i)
            elif i[1] == 2:
                defs.append(i)
            elif i[1] == 3:
                mids.append(i)
            elif i[1] == 4:
                atts.append(i)

        squad = []
        for i in range(len(top_choice[0])):
            if top_choice[0][i] == 1:
                squad.append(goalks[i])
        for i in range(len(top_choice[1])):
            if top_choice[1][i] == 1:
                squad.append(defs[i])
        for i in range(len(top_choice[2])):
            if top_choice[2][i] == 1:
                squad.append(mids[i])
        for i in range(len(top_choice[3])):
            if top_choice[3][i] == 1:
                squad.append(atts[i])

        #amend dataframe to change labels into actual predicted FPL poinnts
        full_squad = pd.DataFrame(squad, columns= ["Player Name", "Element Type", "Team", "Predicted_Points", "Cost"])
        full_squad.insert(loc = 3, column = "Xg_Points", value = "0")
        for index,row in full_squad.iterrows():
            po = full_squad["Predicted_Points"][index]
            if po == 1:
                full_squad.at[index,"Xg_Points"] = "-1"
            elif po == 2:
                full_squad.at[index,"Xg_Points"]  =  "0"
            elif po == 3:
                full_squad.at[index,"Xg_Points"]  = "1"
            elif po == 4:
                full_squad.at[index,"Xg_Points"]  = "2"
            elif po == 5:
                full_squad.at[index,"Xg_Points"]  = "3"
            elif po == 6:
                full_squad.at[index,"Xg_Points"] = "4"
            elif po == 7:
                full_squad.at[index,"Xg_Points"]  = "5"
            elif po == 8:
                full_squad.at[index,"Xg_Points"]  = "6"
            elif po == 9:
                full_squad.at[index,"Xg_Points"] = "7"
            elif po == 10:
                full_squad.at[index,"Xg_Points"]  = "8 - 9"
            elif po == 11:
                full_squad.at[index,"Xg_Points"]  = "10 - 12"
            elif po == 12:
                full_squad.at[index,"Xg_Points"]  = "12 - 14"
            elif po == 13:
                full_squad.at[index,"Xg_Points"]  = "15+"
            elif po >= 14:
                full_squad.at[index,"Xg_Points"]  = str(po+2)
        del full_squad["Predicted_Points"]

        return full_squad, squad

    def initiate(self):
        #Method to initiate the hillcimber
        #@params None
        #@return None

        max_threshold = 100
        generations = 10000
        #mutation rates chosen
        mutation_rates =[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1]
        iterations = 3

        fit_group = []
        totalp = [0]
        totalc = [0]
        final_choices = []
        #make a new squad
        gkp, defs, mids, atks = self.init_hc(max_threshold, self.players)
        for i in range(iterations):
            for mutation_rate in mutation_rates:

                final_choices = []
                #make a copy of the init hill climber so each iteration has the same initial starting squad
                gkp1 = gkp
                mids1 = mids
                defs1 = defs
                atks1 = atks
                fitness_history = []
                #call the hillclimber on the squad
                gkp1, defs1, mids1, atks1, fitness = self.hillclimber(gkp1, defs1, mids1, atks1, self.players, max_threshold, generations, mutation_rate)
                fit_group.append(fitness)
                final_points, final_cost = self.graph(fitness,mutation_rate, self.players, gkp1, defs1, mids1, atks1,generations)
                final_choices.append([gkp1, defs1, mids1, atks1])
                #if squad is optimal, append to optimal squad
                if final_points >= max(totalp):
                    if final_points == max(totalp) and final_cost < max(totalc):
                        continue
                    else:
                        top_choice  = [gkp1,defs1,mids1,atks1]
                        top_cost = final_cost
                        top_points = final_points
                totalp.append(final_points)
                totalc.append(final_cost)
        full_squad, squad = self.project_squad(self.players,top_choice)
        self.full_squad = full_squad
        self.squad = squad
        #top choice is the vector representation of the best squad, used for transfer model
        self.top_choice = top_choice
        print("Squad Selection is Completed")

# In[ ]:


"""Boundary condition implementations.

Boundary Conditions
^^^^^^^^^^^^^^^^^^^

.. autoclass:: PrescribedBoundary
.. autoclass:: DummyBoundary
.. autoclass:: AdiabaticSlipBoundary
"""

__copyright__ = """
Copyright (C) 2020 University of Illinois Board of Trustees
"""

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import numpy as np
from meshmode.dof_array import thaw
from meshmode.mesh import BTAG_ALL, BTAG_NONE  # noqa
# from mirgecom.eos import IdealSingleGas
from grudge.symbolic.primitives import TracePair
from mirgecom.euler import split_conserved, join_conserved

class PrescribedBoundary:
    """Boundary condition prescribes boundary soln with user-specified function.

    .. automethod:: __init__
    .. automethod:: boundary_pair
    """

    def __init__(self, userfunc):
        """Set the boundary function.

        Parameters
        ----------
        userfunc
            User function must take two parameters: time and nodal
            coordinates, and produce the solution at each node.
        """
        self._userfunc = userfunc

    def boundary_pair(
            self, discr, q, btag, eos, t=0.0
    ):
        """Get the interior and exterior solution on the boundary."""
        actx = q[0].array_context

        boundary_discr = discr.discr_from_dd(btag)
        nodes = thaw(actx, boundary_discr.nodes())
        ext_soln = self._userfunc(t, nodes)
        int_soln = discr.project("vol", btag, q)
        return TracePair(btag, interior=int_soln, exterior=ext_soln)

class DummyBoundary:
    """Boundary condition that assigns boundary-adjacent soln as the boundary solution.

    .. automethod:: boundary_pair
    """

    def boundary_pair(
            self, discr, q, btag, eos, t=0.0
    ):
        """Get the interior and exterior solution on the boundary."""
        dir_soln = discr.project("vol", btag, q)
        return TracePair(btag, interior=dir_soln, exterior=dir_soln)

class AdiabaticSlipBoundary:
    r"""Boundary condition implementing inviscid slip boundary.

    a.k.a. Reflective inviscid wall boundary

    This class implements an adiabatic reflective slip boundary given
    by
    $\mathbf{q^{+}} = [\rho^{-}, (\rho{E})^{-}, (\rho\vec{V})^{-}
    - 2((\rho\vec{V})^{-}\cdot\hat{\mathbf{n}}) \hat{\mathbf{n}}]$
    wherein the normal component of velocity at the wall is 0, and
    tangential components are preserved. These perfectly reflecting
    conditions are used by the forward-facing step case in
    [Hesthaven_2008]_, Section 6.6, and correspond to the characteristic
    boundary conditions described in detail in [Poinsot_1992]_.

    .. automethod:: boundary_pair
    """

    def boundary_pair(
            self, discr, q, btag, eos, t=0.0
    ):
        """Get the interior and exterior solution on the boundary.

        The exterior solution is set such that there will be vanishing
        flux through the boundary, preserving mass, momentum (magnitude) and
        energy.
        rho_plus = rho_minus
        v_plus = v_minus - 2 * (v_minus . n_hat) * n_hat
        mom_plus = rho_plus * v_plus
        E_plus = E_minus
        """
        # Grab some boundary-relevant data
        dim = discr.dim
        cv = split_conserved(dim, q)
        actx = cv.mass.array_context

        # Grab a unit normal to the boundary
        nhat = thaw(actx, discr.normal(btag))

        # Get the interior/exterior solns
        int_soln = discr.project("vol", btag, q)
        int_cv = split_conserved(dim, int_soln)

        # Subtract out the 2*wall-normal component
        # of velocity from the velocity at the wall to
        # induce an equal but opposite wall-normal (reflected) wave
        # preserving the tangential component
        mom_normcomp = np.dot(int_cv.momentum, nhat)  # wall-normal component
        wnorm_mom = nhat * mom_normcomp  # wall-normal mom vec
        ext_mom = int_cv.momentum - 2.0 * wnorm_mom  # prescribed ext momentum

        # Form the external boundary solution with the new momentum
        bndry_soln = join_conserved(dim=dim, mass=int_cv.mass,
                                    energy=int_cv.energy,
                                    momentum=ext_mom)

        return TracePair(btag, interior=int_soln, exterior=bndry_soln)


#! /usr/bin/env python

"""
InstrumentData Class -- defines data format, wavelength info, mask geometry

Instruments/masks supported:
NIRISS AMI
GPI, VISIR, NIRC2 removed - too much changed for the JWST NIRISS class
"""

# Standard Imports
import numpy as np
from astropy.io import fits
import os, sys, time
import copy
# Module imports
import synphot
# import stsynphot
# mask geometries, GPI, NIRISS, VISIR supported...
from nrm_analysis.misctools.mask_definitions import NRM_mask_definitions
from nrm_analysis.misctools import utils
from nrm_analysis.misctools import lpl_ianc

um = 1.0e-6

# utility routines for InstrumentData classes

def show_cvsupport_threshold(instr):
    """ Show threshold for where 'splodge' data in CV space contains signal """
    print("InstrumentData: ", "cvsupport_threshold is: ", instr.cvsupport_threshold)
    print("InstrumentData: ", instr.cvsupport_threshold)

def set_cvsupport_threshold(instr, k, v):
    """ Set threshold for where 'splodge' data in CV space contains signal

    Parameters
    ----------
    instr: InstrumentData instance
    thresh: Threshold for the absolute value of the FT(interferogram).
            Normalize abs(CV = FT(a)) for unity peak, and define the support
            of "good" CV when this is above threshold
    """

    instr.cvsupport_threshold[k] = v
    print("InstrumentData: ", "New cvsupport_threshold is: ", instr.cvsupport_threshold)

class NIRISS:
    def __init__(self, filt,
                       objname="obj",
                       src='A0V',
                       chooseholes=None,
                       affine2d=None,
                       bandpass=None,
                       nbadpix=4,
                       usebp=True,
                       firstfew=None,
                       nspecbin=None,
                       **kwargs):
        """
        Initialize NIRISS class

        ARGUMENTS:

        kwargs:
        UTR
        Or just look at the file structure
        Either user has webbpsf and filter file can be read, or...
        chooseholes: None, or e.g. ['B2', 'B4', 'B5', 'B6'] for a four-hole mask
        filt:     Filter name string like "F480M"
        bandpass: None or [(wt,wlen),(wt,wlen),...].  Monochromatic would be e.g. [(1.0, 4.3e-6)]
                  Explicit bandpass arg will replace *all* niriss filter-specific variables with
                  the given bandpass (src, nspecbin, filt), so you can simulate 21cm psfs through
                  something called "F430M". Can also be synphot.spectrum.SourceSpectrum object.
        firstfew: None or the number of slices to truncate input cube to in memory,
                  the latter for fast developmpent
        nbadpix:  Number of good pixels to use when fixing bad pixels DEPRECATED
        usebp:    Convert to usedq during initialization
                  Internally this is changed to sellf.usedq = usebp immediately for code clarity
                  True (default) do not use DQ with DO_NOT_USE flag in input MAST data when
                  fitting data with model.  False: Assume no bad pixels in input
        noise:    standard deviation of noise added to perfect images to enable candid
                  plots without crashing on np.inf limits!  Image assumed to be in (np.float64) dn.
                  Suggested noise: 1e-6.
        src:      source spectral type string e.g. "A0V" OR user-defined synphot.spectrum.SourceSpectrum object
        nspecbin: Number of wavelength bins to use across the bandpass. Replaces deprecated `usespecbin` which
                  set **number of wavelengths to into each bin**, not nbins.

        """

        self.verbose = False
        if "verbose" in kwargs:
            self.verbose = kwargs["verbose"]

        self.noise = None
        if "noise" in kwargs:
            self.noise = kwargs["noise"]
        if "usespecbin" in kwargs: # compatability with previous arg
            # but not really, usespecbin was binning factor, not number of bins
            nspecbin = kwargs["usespecbin"]
        # change how many wavelength bins will be used across the bandpass
        if nspecbin is None:
            nspecbin = 19

        self.lam_bin = nspecbin

        # src can be either a spectral type string or a user-defined synphot spectrum object
        if isinstance(src, synphot.spectrum.SourceSpectrum):
            print("Using user-defined synphot SourceSpectrum")

        if chooseholes:
            print("InstrumentData.NIRISS: ", chooseholes)
        self.chooseholes = chooseholes

        # USEBP is USEDQ in the rest of code - use
        self.usedq = usebp
        print("Fitting omits bad pixels (identified by DO_NOT_USE value in the DQ extension)")
        self.jwst_dqflags() # creates dicts self.bpval, self.bpgroup
        # self.bpexist set True/False if  DQ fits image extension exists/doesn't

        self.firstfew = firstfew
        if firstfew is not None: print("InstrumentData.NIRISS: analysing firstfew={:d} slices".format(firstfew))

        self.objname = objname

        self.filt = filt

        if bandpass is not None:
            print("InstrumentData.NIRISS: OVERRIDING BANDPASS WITH USER-SUPPLIED VALUES.")
            print("\t src, filt, nspecbin parameters will not be used")
            # check type of bandpass. can be synphot spectrum
            # if so, get throughput and wavelength arrays
            if isinstance(bandpass, synphot.spectrum.SpectralElement):
                wl, wt = bandpass._get_arrays(bandpass.waveset)
                self.throughput = np.array((wt,wl)).T
            else:
                self.throughput = np.array(bandpass)  # type simplification
        else:
            filt_spec = utils.get_filt_spec(self.filt)
            src_spec = utils.get_src_spec(src)
            # **NOTE**: As of WebbPSF version 1.0.0 filter is trimmed to where throughput is 10% of peak
            # For consistency with WebbPSF simultions, use trim=0.1
            self.throughput = utils.combine_src_filt(filt_spec,
                                          src_spec,
                                          trim=0.01,
                                          nlambda=nspecbin,
                                          verbose=self.verbose,
                                          plot=False)

        self.lam_c, self.lam_w = utils.get_cw_beta(self.throughput)

        if self.verbose: print("InstrumentData.NIRISS: ", self.filt,
              ": central wavelength {:.4e} microns, ".format(self.lam_c/um), end="")
        if self.verbose: print("InstrumentData.NIRISS: ", "fractional bandpass {:.3f}".format(self.lam_w))

        self.wls = [self.throughput,]

        if self.verbose: print("self.throughput:\n", self.throughput)

        # Wavelength info for NIRISS bands F277W, F380M, F430M, or F480M
        self.wavextension = ([self.lam_c,], [self.lam_w,])
        self.nwav=1 # these are 'slices' if the data is pure imaging integrations -
        #             nwav is old nomenclature from GPI IFU data.  Refactor one day...
        #############################

        # only one NRM on JWST:
        self.telname = "JWST"
        self.instrument = "NIRISS"
        self.arrname = "jwst_g7s6c"  # implaneia mask set with this - unify to short form later
        self.holeshape="hex"
        self.mask = NRM_mask_definitions(maskname=self.arrname, chooseholes=chooseholes,
                                         holeshape=self.holeshape )

        # save affine deformation of pupil object or create a no-deformation object.
        # We apply this when sampling the PSF, not to the pupil geometry.
        # This will set a default Ideal or a measured rotation, for example,
        # and include pixel scale changes due to pupil distortion.
        # Separating detector tilt pixel scale effects from pupil distortion effects is
        # yet to be determined... see comments in Affine class definition.
        # AS AZG 2018 08 15 Ann Arbor
        if affine2d is None:
            self.affine2d = utils.Affine2d(mx=1.0,my=1.0,
                                           sx=0.0,sy=0.0,
                                           xo=0.0,yo=0.0, name="Ideal")
        else:
            self.affine2d = affine2d

        # finding centroid from phase slope only considered cv_phase data
        # when cv_abs data exceeds this cvsupport_threshold.
        # Absolute value of cv data normalized to unity maximum
        # for the threshold application.
        # Data reduction gurus: tweak the threshold value with experience...
        # Gurus: tweak cvsupport with use...
        self.cvsupport_threshold = {"F277W":0.02, "F380M": 0.02, "F430M": 0.02, "F480M": 0.02}
        if self.verbose: show_cvsupport_threshold(self)
        self.threshold = self.cvsupport_threshold[filt]

    def set_pscale(self, pscalex_deg=None, pscaley_deg=None):
        """
        Override pixel scale in header
        """
        if pscalex_deg is not None:
            self.pscalex_deg = pscalex_deg
        if pscaley_deg is not None:
            self.pscaley_deg = pscaley_deg
        self.pscale_mas = 0.5 * (pscalex_deg +  pscaley_deg) * (60*60*1000)
        self.pscale_rad = utils.mas2rad(self.pscale_mas)

    def read_data(self, fn, mode="slice"):
        # mode options are slice or UTR
        # for single slice data, need to read as 3D (1, npix, npix)
        # for utr data, need to read as 3D (ngroup, npix, npix)
        # fix bad pixels using DQ extension and LPL local averaging,
        # but send bad pixel array down to where fringes are fit so they can be ignored.
        # For perfectly noiseless data we add GFaussian zero mean self.noise std dev
        # to imagge data.  Then std devs don't cause plot crashes with limits problems.

        with fits.open(fn, memmap=False, do_not_scale_image_data=True) as fitsfile:
            # use context manager, memmap=False, deepcopy to avoid memory leaks
            scidata = copy.deepcopy(fitsfile[1].data)
            if self.noise is not None: scidata += np.random.normal(0, self.noise, scidata.shape)

            # usually DQ ext in MAST file... make it non-fatal for DQ to be missing
            try:
                bpdata=copy.deepcopy(fitsfile['DQ'].data).astype(np.uint32) # bad pixel extension, forced to uint32
                self.bpexist = True
                dqmask = bpdata & self.bpval["DO_NOT_USE"] == self.bpval["DO_NOT_USE"] #
                del bpdata # free memory

                # True => driver wants to omit using pixels with dqflag raised in fit,
                if self.usedq == True:
                    print('InstrumentData.NIRISS.read_data: will not use flagged DQ pixels in fit')
            except Exception as e:
                print('InstrumentData.NIRISS.read_data: raised exception', e)
                self.bpexist = False
                dqmask = np.zeros(scidata.shape, dtype=np.uint32) # so it doesn't break if issues with DQ data

            if scidata.ndim == 3:  #len(scidata.shape)==3:
                print("read_data() input: 3D cube")

                # Truncate all but the first few slices od data and DQ array for rapid development
                if self.firstfew is not None:
                    if scidata.shape[0] > self.firstfew:
                        scidata = scidata[:self.firstfew, :, :]
                        dqmask = dqmask[:self.firstfew, :, :]
                # 'nwav' name (historical) is actually number of data slices in the 3Dimage cube
                self.nwav=scidata.shape[0]
                [self.wls.append(self.wls[0]) for f in range(self.nwav-1)]

            elif len(scidata.shape)==2: # 'cast' 2d array to 3d with shape[0]=1
                print("'InstrumentData.NIRISS.read_data: 2D data array converting to 3D one-slice cube")
                scidata = np.array([scidata,])
                dqmask = np.array([dqmask,])
            else:
                sys.exit("InstrumentData.NIRISS.read_data: invalid data dimensions for NIRISS. \nShould have dimensionality of 2 or 3.")

            # refpix removal by trimming
            scidata = scidata[:,4:, :] # [all slices, imaxis[0], imaxis[1]]
            print('\tRefpix-trimmed scidata:', scidata.shape)
            #### fix pix using bad pixel map - runs now.  Need to sanity-check.
            if self.bpexist:
                # refpix removal by trimming to match image trim
                dqmask = dqmask[:,4:, :]     # dqmask bool array to match image trimmed shape
                print('\tRefpix-trimmed dqmask: ', dqmask.shape)

            prihdr=fitsfile[0].header
            scihdr=fitsfile[1].header
            # MAST header or similar kwds info for oifits writer:
            self.updatewithheaderinfo(prihdr, scihdr)

            # Directory name into which to write txt observables & optional fits diagnostic files
            # The input fits image or cube of images file rootname is used to create the output
            # text&fits dir, using the data file's root name as the directory name: for example,
            # /abc/.../imdir/xyz_calints.fits  results in a directory /abc/.../imdir/xyz_calints/
            self.rootfn =  fn.split('/')[-1].replace('.fits', '')
        return prihdr, scihdr, scidata, dqmask

    def cdmatrix_to_sky(self, vec, cd11, cd12, cd21, cd22):
        """ use the global header values explicitly, for clarity
            vec is 2d, units of pixels
            cdij 4 scalars, conceptually 2x2 array in units degrees/pixel
        """
        return np.array((cd11*vec[0] + cd12*vec[1], cd21*vec[0] + cd22*vec[1]))

    def degrees_per_pixel(self, hdr):
        """
        input: hdr:  fits data file's header with or without CDELT1, CDELT2 (degrees per pixel)
        returns: cdelt1, cdelt2: tuple, degrees per pixel along axes 1, 2
                         EITHER: read from header CDELT[12] keywords
                             OR: calculated using CD matrix (Jacobian of RA-TAN, DEC-TAN degrees
                                 to pixel directions 1,2.  No deformation included in this routine,
                                 but the CD matric includes non-linear field distortion.
                                 No pupil distortion or rotation here.
                        MISSING: If keywords are missing default hardcoded cdelts are returned.
                                 The exact algorithm may substitute this later.
                                 Below seems good to ~5th significant figure when compared to
                                 cdelts header values prior to replacement by cd matrix approach.
            N.D. at stsci 11 Mar 20212

            We start in Level 1 with the PC matrix and CDELT.
            CDELTs come from the SIAF.

            The PC matrix is computed from the roll angle, V3YANG and the parity.
            The code is here
            https://github.com/spacetelescope/jwst/blob/master/jwst/assign_wcs/util.py#L153

            In the level 2 imaging pipeline, assign_wcs adds the distortion to the files.
            At the end it computes an approximation of the entire distortion transformation
            by fitting a polynomial. This approximated distortion is represented as SIP
            polynomials in the FITS headers.

            Because SIP, by definition, uses a CD matrix, the PC + CDELT are replaced by CD.

            How to get CDELTs back?

            I think once the rotation, skew and scale are in the CD matrix it's very hard to
            disentangle them. The best way IMO is to calculate the local scale using three
            point difference. There is a function in jwst that does this.

            Using a NIRISS image as an example:

            from jwst.assign_wcs import util
            from jwst import datamodels

            im=datamodels.open('niriss_image_assign_wcs.fits')

            util.compute_scale(im.meta.wcs, (im.meta.wcsinfo.ra_ref, im.meta.wcsinfo.dec_ref))

            1.823336635353374e-05

            The function returns a constant scale. Is this sufficient for what you need or
            do you need scales and sheer along each axis? The code in util.compute_scale can
            help with figuring out how to get scales along each axis.

            I hope this answers your question.
        """

        if 'CD1_1' in hdr.keys() and 'CD1_2' in hdr.keys() and  \
             'CD2_1' in hdr.keys() and 'CD2_2' in hdr.keys():
            cd11 = hdr['CD1_1']
            cd12 = hdr['CD1_2']
            cd21 = hdr['CD2_1']
            cd22 = hdr['CD2_2']
            # Create unit vectors in detector pixel X and Y directions, units: detector pixels
            dxpix  =  np.array((1.0, 0.0)) # axis 1 step
            dypix  =  np.array((0.0, 1.0)) # axis 2 step
            # transform pixel x and y steps to RA-tan, Dec-tan degrees
            dxsky = self.cdmatrix_to_sky(dxpix, cd11, cd12, cd21, cd22)
            dysky = self.cdmatrix_to_sky(dypix, cd11, cd12, cd21, cd22)
            print("Used CD matrix for pixel scales")
            return np.linalg.norm(dxsky, ord=2), np.linalg.norm(dysky, ord=2)
        elif 'CDELT1' in hdr.keys() and 'CDELT2' in hdr.keys():
            return hdr['CDELT1'], hdr['CDELT2']
            print("Used CDDELT[12] for pixel scales")
        else:
            print('InstrumentData.NIRISS: Warning: NIRISS pixel scales not in header.  Using 65.6 mas in deg/pix')
            return 65.6/(60.0*60.0*1000), 65.6/(60.0*60.0*1000)

    def updatewithheaderinfo(self, ph, sh):
        """ input: primary header, science header MAST"""

        # The info4oif_dict will get pickled to disk when we write txt files of results.
        # That way we don't drag in objects like InstrumentData into code that reads text results
        # and writes oifits files - a simple built-in dictionary is the only object used in this transfer.
        info4oif_dict = {}
        info4oif_dict['telname'] = self.telname

        info4oif_dict['filt'] = self.filt
        info4oif_dict['lam_c'] = self.lam_c
        info4oif_dict['lam_w'] = self.lam_w
        info4oif_dict['lam_bin'] = self.lam_bin

        # Target information - 5/21 targname UNKNOWN in nis019 rehearsal data
        # Name in the proposal always non-trivial, targname still UNKNOWN...:
        if ph["TARGNAME"] == 'UNKNOWN': objname = ph['TARGPROP']
        else: objname = ph['TARGNAME'] # allegedly apt name for archive, standard form
        #
        # if target name has confusing-to-astroquery dash
        self.objname =  objname.replace('-', ' '); info4oif_dict['objname'] = self.objname
        # AB Dor, ab dor, AB DOR,  ab  dor are all acceptable.
        #
        self.ra = ph["TARG_RA"]; info4oif_dict['ra'] = self.ra
        self.dec = ph["TARG_DEC"]; info4oif_dict['dec'] = self.dec

        # / axis 1 DS9 coordinate of the reference pixel (always POS1)
        # / axis 2 DS9 coordinate of the reference pixel (always POS1)
        self.crpix1 = sh["CRPIX1"]; info4oif_dict['crpix1'] = self.crpix1
        self.crpix2 = sh["CRPIX2"]; info4oif_dict['crpix2'] = self.crpix2
        # need Paul Goudfrooij's table for actual crval[1,2] for true pointing to detector pixel coords (DS9)

        self.instrument = ph["INSTRUME"]; info4oif_dict['instrument'] = self.instrument
        self.pupil =  ph["PUPIL"]; info4oif_dict['pupil'] = self.pupil
        # "ImPlaneIA internal mask name" - oifwriter looks for 'mask'...
        self.arrname = "jwst_g7s6c"  # implaneia internal name - historical
        info4oif_dict['arrname'] = 'g7s6' # for oif
        info4oif_dict['mask'] = info4oif_dict['arrname']  # Soulain mask goes into oif arrname

        # if data was generated on the average pixel scale of the header
        # then this is the right value that gets read in, and used in fringe fitting
        pscalex_deg, pscaley_deg = self.degrees_per_pixel(sh)
        #
        info4oif_dict['pscalex_deg'] = pscalex_deg
        info4oif_dict['pscaley_deg'] = pscaley_deg
        # Whatever we did set is averaged for isotropic pixel scale here
        self.pscale_mas = 0.5 * (pscalex_deg + pscaley_deg) * (60*60*1000); \
        info4oif_dict['pscale_mas'] = self.pscale_mas
        self.pscale_rad = utils.mas2rad(self.pscale_mas); info4oif_dict['pscale_rad'] = self.pscale_rad

        self.mask = NRM_mask_definitions(maskname=self.arrname, chooseholes=self.chooseholes,
                                         holeshape=self.holeshape) # for STAtions x y in oifs

        self.date = ph["DATE-OBS"] + "T" + ph["TIME-OBS"]; info4oif_dict['date'] = self.date
        datestr = ph["DATE-OBS"]
        self.year = datestr[:4]; info4oif_dict['year'] = self.year
        self.month = datestr[5:7]; info4oif_dict['month'] = self.month
        self.day = datestr[8:10]; info4oif_dict['day'] = self.day
        self.parangh= sh["ROLL_REF"]; info4oif_dict['parangh'] = self.parangh
        self.pa = sh["PA_V3"]; info4oif_dict['pa'] = self.pa
        self.vparity = sh["VPARITY"]; info4oif_dict['vparity'] = self.vparity

        # An INTegration is NGROUPS "frames", not relevant here but context info.
        # 2d => "cal" file combines all INTegrations (ramps)
        # 3d=> "calints" file is a cube of all INTegrations (ramps)
        if sh["NAXIS"] == 2:
            # all INTegrations or 'ramps'
            self.itime = ph["EFFINTTM"] * ph["NINTS"]; info4oif_dict['itime'] = self.itime
        elif sh["NAXIS"] == 3:
            # each slice is one INTegration or 'ramp'
            self.itime = ph["EFFINTTM"]; info4oif_dict['itime'] = self.itime

        np.set_printoptions(precision=5, suppress=True, linewidth=160,
                            formatter={'float': lambda x: "%10.5f," % x})
        self.v3i_yang = sh['V3I_YANG']  # Angle from V3 axis to Ideal y axis (deg)
        # rotate mask hole center coords by PAV3 # RAC 2021
        ctrs_sky = self.mast2sky()
        oifctrs = np.zeros(self.mask.ctrs.shape)
        oifctrs[:,0] = ctrs_sky[:,1].copy() * -1
        oifctrs[:,1] = ctrs_sky[:,0].copy() * -1
        info4oif_dict['ctrs_eqt'] = oifctrs # mask centers rotated by PAV3 (equatorial coords)
        info4oif_dict['ctrs_inst'] = self.mask.ctrs # as-built instrument mask centers
        info4oif_dict['hdia'] = self.mask.hdia
        info4oif_dict['nslices'] = self.nwav # nwav: number of image slices or IFU cube slices - AMI is imager
        self.info4oif_dict = info4oif_dict # save it when writing extracted observables txt

    # rather than calling InstrumentData in the niriss example just to reset just call this routine
    def reset_nwav(self, nwav):
        print("InstrumentData.NIRISS: ", "Resetting InstrumentData instantiation's nwave to", nwav)
        self.nwav = nwav

    def jwst_dqflags(self):
        """
            dqdata is a 2d (32-bit U?)INT array from the DQ extension of the input file.
            We ignore all data with a non-zero DQ flag.  I copied all values from a 7.5 build jwst...
            but we ignore any non-zero flag meaning, and ignore the pixel in fringe-fitting
            The refpix are non-zero DQ, btw...
            I changed "pixel" to self.pbval and "group" to self.bpgroup. We may use these later,
            so here they are but initially we just discriminate between good (zero value) and non-good.
        """

        """ JWST Data Quality Flags
            The definitions are documented in the JWST RTD:
            https://jwst-pipeline.readthedocs.io/en/latest/jwst/references_general/references_general.html#data-quality-flags
        """
        """ JWST Data Quality Flags
        The definitions are documented in the JWST RTD:
        https://jwst-pipeline.readthedocs.io/en/latest/jwst/references_general/references_general.html#data-quality-flags
        Implementation
        -------------
        The flags are implemented as "bit flags": Each flag is assigned a bit position
        in a byte, or multi-byte word, of memory. If that bit is set, the flag assigned
        to that bit is interpreted as being set or active.
        The data structure that stores bit flags is just the standard Python `int`,
        which provides 32 bits. Bits of an integer are most easily referred to using
        the formula `2**bit_number` where `bit_number` is the 0-index bit of interest.
        2**n is gauche but not everyone loves 1<<n

        Rachel uses:
        from jwst.datamodels import dqflags
        DO_NOT_USE = dqflags.pixel["DO_NOT_USE"]
        dqmask = pxdq0 & DO_NOT_USE == DO_NOT_USE
        pxdq = np.where(dqmask, pxdq0, 0)

        """

        # Pixel-specific flags
        self.bpval = {
                 'GOOD':             0,      # No bits set, all is good
                 'DO_NOT_USE':       2**0,   # Bad pixel. Do not use.
                 'SATURATED':        2**1,   # Pixel saturated during exposure
                 'JUMP_DET':         2**2,   # Jump detected during exposure
                 'DROPOUT':          2**3,   # Data lost in transmission
                 'OUTLIER':          2**4,   # Flagged by outlier detection. Was RESERVED_1
                 'RESERVED_2':       2**5,   #
                 'RESERVED_3':       2**6,   #
                 'RESERVED_4':       2**7,   #
                 'UNRELIABLE_ERROR': 2**8,   # Uncertainty exceeds quoted error
                 'NON_SCIENCE':      2**9,   # Pixel not on science portion of detector
                 'DEAD':             2**10,  # Dead pixel
                 'HOT':              2**11,  # Hot pixel
                 'WARM':             2**12,  # Warm pixel
                 'LOW_QE':           2**13,  # Low quantum efficiency
                 'RC':               2**14,  # RC pixel
                 'TELEGRAPH':        2**15,  # Telegraph pixel
                 'NONLINEAR':        2**16,  # Pixel highly nonlinear
                 'BAD_REF_PIXEL':    2**17,  # Reference pixel cannot be used
                 'NO_FLAT_FIELD':    2**18,  # Flat field cannot be measured
                 'NO_GAIN_VALUE':    2**19,  # Gain cannot be measured
                 'NO_LIN_CORR':      2**20,  # Linearity correction not available
                 'NO_SAT_CHECK':     2**21,  # Saturation check not available
                 'UNRELIABLE_BIAS':  2**22,  # Bias variance large
                 'UNRELIABLE_DARK':  2**23,  # Dark variance large
                 'UNRELIABLE_SLOPE': 2**24,  # Slope variance large (i.e., noisy pixel)
                 'UNRELIABLE_FLAT':  2**25,  # Flat variance large
                 'OPEN':             2**26,  # Open pixel (counts move to adjacent pixels)
                 'ADJ_OPEN':         2**27,  # Adjacent to open pixel
                 'UNRELIABLE_RESET': 2**28,  # Sensitive to reset anomaly
                 'MSA_FAILED_OPEN':  2**29,  # Pixel sees light from failed-open shutter
                 'OTHER_BAD_PIXEL':  2**30,  # A catch-all flag
                 'REFERENCE_PIXEL':  2**31,  # Pixel is a reference pixel
        }

        # Group-specific flags. Once groups are combined, these flags
        # are equivalent to the pixel-specific flags.
        self.bpgroup = {
                 'GOOD':       self.bpval['GOOD'],
                 'DO_NOT_USE': self.bpval['DO_NOT_USE'],
                 'SATURATED':  self.bpval['SATURATED'],
                 'JUMP_DET':   self.bpval['JUMP_DET'],
                 'DROPOUT':    self.bpval['DROPOUT'],
        }

    def mast2sky(self):
        """
        Rotate hole center coordinates:
            Clockwise by the V3 position angle - V3I_YANG from north in degrees if VPARITY = -1
            Counterclockwise by the V3 position angle - V3I_YANG from north in degrees if VPARITY = 1
        Hole center coords are in the V2, V3 plane in meters.
        Return rotated coordinates to be put in info4oif_dict.
        implane2oifits.ObservablesFromText uses these to calculate baselines.
        """
        pa = self.pa
        mask_ctrs = copy.deepcopy(self.mask.ctrs)
        # rotate by an extra 90 degrees (RAC 9/21)
        # these coords are just used to orient output in OIFITS files
        # NOT used for the fringe fitting itself
        mask_ctrs = utils.rotate2dccw(mask_ctrs,np.pi/2.)
        vpar = self.vparity # Relative sense of rotation between Ideal xy and V2V3
        v3iyang = self.v3i_yang
        rot_ang = pa - v3iyang # subject to change!

        if pa != 0.0:
            # Using rotate2sccw, which rotates **vectors** CCW in a fixed coordinate system,
            # so to rotate coord system CW instead of the vector, reverse sign of rotation angle.  Double-check comment
            if vpar == -1:
                # rotate clockwise  <rotate coords clockwise?>
                ctrs_rot = utils.rotate2dccw(mask_ctrs, np.deg2rad(-rot_ang))
                print(f'InstrumentData.mast2sky: Rotating mask hole centers clockwise by {rot_ang:.3f} degrees')
            else:
                # counterclockwise  <rotate coords counterclockwise?>
                ctrs_rot = utils.rotate2dccw(mask_ctrs, np.deg2rad(rot_ang))
                print('InstrumentData.mast2sky: Rotating mask hole centers counterclockwise by {rot_ang:.3f} degrees')
        else:
            ctrs_rot = mask_ctrs
        return ctrs_rot


import argparse
import glob
import itertools
from nuosc.tools import cache, progress
import numpy
import ROOT
import math
import simplot.io
import nuosc.model
from simplot import histogram
import matplotlib.pyplot as plt
import os
import scipy.stats
import matplotlib
from nuosc.model import constants, runtime
import collections
import re
import scipy.interpolate

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

GenEvFile = collections.namedtuple("GenEvFile", ["polarity", "ndid", "r", "z", "n", "orientation"])
BeamFile = collections.namedtuple("BeamFile", ["polarity"])

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

class FileNameParser:
    def __init__(self, patterns):
        self._flist = _expandfilelist(patterns)

    def __call__(self):
        r = None
        fname = self._flist[0]
        for f in [self._parse_genev, self._parse_beam]:
            r = f(fname)
            if r is not None:
                break
        #check for success
        if r is None:
            raise ValueError("cannot parse ", fname)
        return r

    def _parse_genev(self, fname):
        pat = "genev_(.*?)_(.*?)_(.*?)_.(.*?)_z(.*?)_(.)_(.*?).root"
        match = re.search(pat, fname)
        if not match:
            return None
        polarity = match.group(1)
        det = match.group(2)
        ndid = nuosc.model.constants.DetectorId.tostring(det)
        geo = match.group(3)
        r = float(int(match.group(4))) / 100.0
        z = float(int(match.group(5))) / 100.0
        orientation = match.group(6)
        n = match.group(7)
        return GenEvFile(polarity=polarity,
                  ndid=ndid,
                  r=r,
                  z=z,
                  n=n,
                  orientation=orientation,
                  )

    def _parse_beam(self, fname):
        pat = "merged_beamfiles_(.*?).root"
        match = re.search(pat, fname)
        if not match:
            return None
        polarity = match.group(1)
        return BeamFile(polarity=polarity)

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

class Record:
    pdg = "pdg"
    reactioncode = "reactioncode"
    enu = "enu"
    x = "x"
    y = "y"
    z = "z"
    weight = "weight"
    ndid = "ndid"

    ALL = (pdg, reactioncode, enu, x, y, z, weight)

    _axis_labels = {pdg : "PDG code",
                    reactioncode : "NEUT reaction code",
                    enu : r"$E_{\nu}$ [GeV]",
                    x : "$x$ [m]",
                    y : "$y$ [m]",
                    z : "$z$ [m]",
                    weight : "event weight",
                    ndid : "near detector ID",
                    }

    @classmethod
    def label(cls, field):
        return cls._axis_labels[field]

def _expandfilelist(filelist):
    return list(itertools.chain.from_iterable(map(glob.glob, filelist)))

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

class NeutEventReader:

    def __call__(self, event):
        x, y, z = self._pos(event)
        return [event.StdHepPdg[0],
                event.NEneutmode,
                self._enu(event),
                x,
                y,
                z,
                1.0,
                0.0,
                ]

    def _pos(self, event):
        r = event.EvtVtx
        return (r[0],
                r[1],
                r[2],
                )

    def _convert2Darray(self, arr, nrows, ncols):
        out = numpy.zeros(shape=(nrows,ncols))
        for i in xrange(nrows):
            for j in xrange(ncols):
                out[i, j] = arr[(i * ncols) + j]
        return out

    def _enu(self, event):
        arr = self._convert2Darray(event.StdHepP4, event.StdHepN, 4)
        return arr[0][3]

    def indices(self):
        return [Record.pdg,
                Record.reactioncode,
                Record.enu,
                Record.x,
                Record.y,
                Record.z,
                Record.weight,
                Record.ndid,
                ]

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

class FluxEventReader:

    def __call__(self, event):
        return [event.pdg,
                0.0,
                event.enu,
                event.xnu,
                event.ynu,
                0.0,
                event.norm,
                event.ndid,
                ]

    def indices(self):
        return [Record.pdg,
                Record.reactioncode,
                Record.enu,
                Record.x,
                Record.y,
                Record.z,
                Record.weight,
                Record.ndid,
                ]

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

class Reader(object):
    def __init__(self, filelist, treename, func, prescale=1):
        self._filenames = _expandfilelist(filelist)
        self._func = func
        self._treename = treename
        self._prescale = prescale

    def read(self):
        indices = self._func.indices()
        uniquestr = "_".join(indices
                             + self._filenames
                             + [type(self._func).__name__,
                                str(self._prescale),
                                ]
                             )
        cachetool = cache.CacheNumpy(uniquestr, prefix="tmp_verify_genev")
        if cachetool.exists():
            data = cachetool.read()
        else:
            data = self._load_data()
            cachetool.write(data)
        indices = dict(zip(indices, xrange(len(indices))))
        return indices, data

    def _load_data(self):
        data = []
        for fname in self._filenames:
            tfile = ROOT.TFile(fname, "READ")
            treename = self._treename
            tree = tfile.Get(treename)
            if not tree:
                raise Exception("Input file does not contain tree with name", treename, fname)
            for event in tree:
                data.append(self._func(event))
        print len(data), data[0]
        data = numpy.array(data)
        return data

    def _iterprescale(self, tree, prescale):
        n = tree.GetEntries()
        for i in xrange(n):
            if i % prescale == 0:
                tree.GetEntry(i)
                yield tree

    def _load_data(self):
        data = []
        for fname in self._filenames:
            tfile = ROOT.TFile(fname, "READ")
            treename = self._treename
            tree = tfile.Get(treename)
            if not tree:
                raise Exception("Input file does not contain tree with name", treename, fname)
            prescale = self._prescale
            iterevents = progress.printProgress("Reading " + treename, tree.GetEntries()/prescale, self._iterprescale(tree, prescale))
            for event in iterevents:
                data.append(self._func(event))
        print len(data), data[0]
        data = numpy.array(data)
        return data

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

class NeutFileReader(Reader):
    def __init__(self, filelist):
        super(NeutFileReader, self).__init__(filelist, "nRooTracker", NeutEventReader(), prescale=1)

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

class FluxFileReader(Reader):
    def __init__(self, filelist):
        super(FluxFileReader, self).__init__(filelist, "flux", FluxEventReader(), prescale=20)

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

class PlotInteractionRate:
    def __init__(self, outdir, indices, data, scale_to_m=1.0, namemod=None, numbins=25):
        self._scale_to_m = scale_to_m
        self._outdir = outdir
        self._data = data
        self._indices = indices
        self._namemod = namemod
        self._numbins = numbins

    def __iter__(self):
        if len(self._data):
            for p in self._plot_xy_hist():
                yield p
            for p in self._plot_xz_hist():
                yield p
#             for p in self._plot_histograms():
#                 yield p

    def _plot_histograms(self):
        for var in Record.ALL:
            fig = plt.figure()
            d = self._data[:, self._indices[var]]
            weights = self._data[:, self._indices[Record.weight]]
            y, xbinning = numpy.histogram(d, weights=weights)
            ax = fig.add_subplot(1, 1, 1)
            histogram.plot_hist_points(ax, xbinning, y)
            ax.set_xlabel(Record.label(var))
            ax.set_ylabel("N events")
            name = os.sep.join((self._outdir, "hist_1d", var))
            if self._namemod is not None:
                name += "_" + self._namemod
            yield name, fig
            plt.close(fig)

    def _plot_xy_hist(self):
        fig = plt.figure()
        name = os.sep.join((self._outdir, "hist_2d", "xy"))
        if self._namemod is not None:
            name += "_" + self._namemod
        ax = fig.add_subplot(1, 1, 1)
        ax.set_xlabel(Record.label(Record.x))
        ax.set_ylabel(Record.label(Record.y))
        #Get data
        dx = self._data[:, self._indices[Record.x]]
        dy = self._data[:, self._indices[Record.y]]
        #scale dx and dy
        dx *= self._scale_to_m
        dy *= self._scale_to_m
        weights = self._data[:, self._indices[Record.weight]]
        #Plot data
        xmin = -10.0
        xmax = 10.0
        prange = ((xmin, xmax), (xmin, xmax))
        ret = ax.hist2d(dx, dy, bins=self._numbins, weights=weights, cmap=matplotlib.cm.get_cmap("hot"), range=prange)
        img = ret[-1]
        fig.colorbar(img)
        yield name, fig

    def _plot_xz_hist(self):
        fig = plt.figure()
        name = os.sep.join((self._outdir, "hist_2d", "xz"))
        if self._namemod is not None:
            name += "_" + self._namemod
        ax = fig.add_subplot(1, 1, 1)
        ax.set_xlabel(Record.label(Record.x))
        ax.set_ylabel(Record.label(Record.z))
        #Get data
        dx = self._data[:, self._indices[Record.x]]
        dz = self._data[:, self._indices[Record.z]]
        #scale dx and dy
        dx *= self._scale_to_m
        dz *= self._scale_to_m
        weights = self._data[:, self._indices[Record.weight]]
        #Plot data
        #xmin = -4.0
        #xmax = 4.0
        #prange = ((xmin, xmax), (xmin, xmax))
        ret = ax.hist2d(dx, dz, bins=self._numbins, weights=weights, cmap=matplotlib.cm.get_cmap("hot"), )#range=prange)
        img = ret[-1]
        fig.colorbar(img)
        yield name, fig

#     def _plot_xy_kernel(self):
#         fig = plt.figure()
#         ax = fig.add_subplot(1, 1, 1)
#         ax.set_xlabel(Record.label(Record.x))
#         ax.set_ylabel(Record.label(Record.y))
#         #Get data
#         dx = self._data[:, self._indices[Record.x]]
#         dy = self._data[:, self._indices[Record.y]]
#         #Plot histogram
#         heatmap, xedges, yedges = numpy.histogram2d(dx, dy, bins=20)
#         extent = [min(xedges), max(xedges), min(yedges), max(yedges)]
#         ax.imshow(heatmap, extent=extent, interpolation="nearest")
#         dataset = numpy.row_stack((dx, dy))
#         x = numpy.linspace(min(dx), max(dx))
#         y = numpy.linspace(min(dy), max(dy))
#         X, Y = numpy.meshgrid(x, y)
#         kernel = scipy.stats.gaussian_kde(dataset)
#         Z = kernel(numpy.row_stack((X.ravel(), Y.ravel()))).reshape(X.shape)
#         ax.contour(X, Y, Z, color="black", linecolor="black")
#         name = os.sep.join(("hist_2d", "xy"))
#         yield name, fig

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

class Flux:
    def __init__(self, data, weights):
        #bin the data
        nbins = 100.0
        enurange = (0.0, 5.0)
        hist, edge = numpy.histogram(data, bins=nbins, range=enurange, density=True, weights=weights)
        y = hist
        x = (edge[1:] + edge[:-1] ) / 2.0
        self.f = scipy.interpolate.interp1d(x, y)

    def __call__(self, x):
        return self.f(x)

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

class Xsec:
    def __init__(self):
        pass

    def __call__(self, x):
        return 1.0

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

class FluxIntXsec:
    def __init__(self, flux, xsec):
        enu = numpy.linspace(start=0.01, stop=5.0, num=1000)
        ff = numpy.vectorize(flux)
        fa = ff(enu)
        xf = numpy.vectorize(xsec)
        xa = xf(enu)
        avg = numpy.average(xa, weights=fa)
        self._result = avg

    def __call__(self):
        return self._result

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

class Ratio(object):
    def __init__(self, numerator, denominator):
        self._n = numerator
        self._d = denominator
        try:
            self._r = float(self._n) / float(self._d)
        except ZeroDivisionError:
            self._r = 0.0

    def __div__(self, rhs):
        return Ratio(self, rhs)

    def __sub__(self, rhs):
        return Difference(self, rhs)

    def __float__(self):
        return self._r

    def __str__(self):
        #return "{0:.3e} +- {1:.3e}".format(self._r, self.error())
        return "{0:.3f} +- {1:.3f}".format(self._r, self.error())
#         return "\n".join(("Num="+str(self._n),
#                          "Den="+str(self._d),
#                          "R="+str(self._r),
#                          ))

    def error(self):
        err = 0.0
        try:
            n = float(self._n)
            d = float(self._d)
            ne = self._n.error()
            de = self._d.error()
            err = self._r * math.sqrt((ne/n)**2 + (de/d)**2)
        except AttributeError:
            #assume binomial error
            eff = self._r
            err = 0.0
            if self._d > 0:
                var = eff * (1.0 - eff) / float(self._n)
            err = math.sqrt(var)
        return err

class Difference(Ratio):
    def __init__(self, numerator, denominator):
        self._n = numerator
        self._d = denominator
        self._r = float(self._n) - float(self._d)

    def error(self):
        err = 0.0
        n = float(self._n)
        d = float(self._d)
        ne = self._n.error()
        de = self._d.error()
        err = math.sqrt((ne)**2 + (de)**2)
        return err

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

class PlotFluxIntXsec:
    def __init__(self, neutindices, neutdata, fluxindices, fluxdata, scale_to_m = 1.0/100.0):
        self._neut_indices = neutindices
        self._neut_data = neutdata
        self._fluxindices = fluxindices
        self._fluxdata = fluxdata
        self._scale_to_m = scale_to_m

    def __iter__(self):
        rad = 4.0
        outer_xmax = rad / math.sqrt(2.0)
        outer_ymax = outer_xmax
        outer_xmin = -outer_xmax
        outer_ymin = -outer_ymax
        #get fiducial flux
        fvflux = _spacecut(self._fluxindices, self._fluxdata, outer_xmin, outer_xmax, outer_ymin, outer_ymax, scale_to_m=self._scale_to_m)
        fvneut = _spacecut(self._neut_indices, self._neut_data, outer_xmin, outer_xmax, outer_ymin, outer_ymax)
        #divide square into quadrants
        rad = 4.0
        outer_xmax = rad / math.sqrt(2.0)
        outer_ymax = outer_xmax
        outer_xmin = -outer_xmax
        outer_ymin = -outer_ymax
        for name, xmin, xmax, ymin, ymax in [("(+x, +y)", 0.0, outer_xmax, 0.0, outer_ymax),
                                             ("(+x, -y)", 0.0, outer_xmax, outer_ymin, 0.0),
                                             ("(-x, +y)", outer_xmin, 0.0, 0.0, outer_ymax),
                                             ("(-x, -y)",outer_xmin, 0.0, outer_ymin, 0.0),

                                       ]:
            #print "Plotting ", name
            for p in self._plot(xmin, xmax, ymin, ymax, fvflux, fvneut, self._fluxindices, self._neut_indices):
                yield p
        return

    def _plot(self, xmin, xmax, ymin, ymax, fluxdata, neut_data, fluxindices, neut_indices):
        fd = _spacecut(fluxindices, fluxdata, xmin, xmax, ymin, ymax, scale_to_m=self._scale_to_m)
        gd = _spacecut(neut_indices, neut_data, xmin, xmax, ymin, ymax)
#         #calc flux integrated xsec
#         flux = Flux(fd[:, Record.enu], fd[:, Record.weight])
#         xsec = Xsec()
#         fluxintxsec = FluxIntXsec(flux, xsec)
        #yield figname, fig
        expected_frac = Ratio(self._sumflux(fluxindices, fd), self._sumflux(fluxindices, fluxdata))
        observed_frac = Ratio(len(gd), len(neut_data))
        diff = observed_frac - expected_frac
        #print "expected=", str(expected_frac), ", observed=", str(observed_frac),", diff=",(diff)
        l = [(100.0*float(v), 100.0*v.error()) for v in (expected_frac, observed_frac, diff)]
        for t in l:
            print t,","
        yield None

    def _sumflux(self, indices, fd):
        return numpy.sum(fd[:, indices[Record.weight]])

def _spacecut(indices, fd, xmin, xmax, ymin, ymax, scale_to_m=1.0):
        xmin, xmax, ymin, ymax = [[v / scale_to_m] for v in [xmin, xmax, ymin, ymax]]
        #x-cut
        ix = indices[Record.x]
        fd = fd[fd[:, ix] > xmin]
        fd = fd[fd[:, ix] < xmax]
        #y cut
        iy = indices[Record.y]
        fd = fd[fd[:, iy] > ymin]
        fd = fd[fd[:, iy] < ymax]
        return fd

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

def plot_interaction_rate(args):
    parser = FileNameParser(args.patterns)
    f = parser()
    reader = NeutFileReader(args.patterns)
    indices, data = reader.read()
    out = simplot.io.FigureWriter()
    name = "_".join([f.polarity, f.ndid])
    plotter = PlotInteractionRate("evgen", indices, data,
                                  namemod=name,
                                  numbins=100,
                                  )
    out(plotter)
    return

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

def plot_flux(args):
    ntuple = nuosc.model.beam.FluxNtuple(runtime.getcontext().beamcontext)
    reader = FluxFileReader([ntuple.filename()])
    indices, data = reader.read()
    DetectorId = nuosc.model.beam.constants.DetectorId
    out = simplot.io.FigureWriter()
    for polarity in nuosc.model.constants.Polarity.ALL:
        for ndid in DetectorId.ALL:
            #filter out particular near detector
            indid = DetectorId.toint(ndid)
            fluxdata = data[data[:, indices[Record.ndid]] == indid]
            if len(fluxdata) > 0:
                if polarity > 0:
                    fluxdata = fluxdata[fluxdata[:, indices[Record.pdg]] > 0]
                else:
                    fluxdata = fluxdata[fluxdata[:, indices[Record.pdg]] < 0]
            detstr = DetectorId.tostring(ndid)
            cm_to_m = 1.0 / 100.0
            name = "_".join((detstr, str(polarity)))
            plotter = PlotInteractionRate("flux", indices, fluxdata,
                                          scale_to_m=cm_to_m,
                                          namemod=name,
                                          numbins=10,
                                          )
            out(plotter)

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

def plot_flux_int_xsec(args):
    parser = FileNameParser(args.patterns)
    f = parser()
    #load neut data
    reader = NeutFileReader(args.patterns)
    neutindices, neutdata = reader.read()
    #load flux data
    ntuple = nuosc.model.beam.FluxNtuple(runtime.getcontext().beamcontext)
    reader = FluxFileReader([ntuple.filename()])
    fluxindices, fluxdata = reader.read()
    DetectorId = nuosc.model.beam.constants.DetectorId
    ndid = DetectorId.toint(f.ndid)
    #filter out particular near detector
    fluxdata = fluxdata[fluxdata[:, fluxindices[Record.ndid]] == ndid]
    #make plot
    out = simplot.io.FigureWriter()
    cm_to_m = 1.0 / 100.0
    #cut on flux energy
    plotter = PlotFluxIntXsec(neutindices, neutdata, fluxindices, fluxdata,
                              scale_to_m=cm_to_m,
                              )
    #out(plotter)
    for p in plotter:
        pass

def testcut(indices, d, scale_to_m):
        rad = 3.0
        outer_xmax = rad / math.sqrt(2.0)
        outer_ymax = outer_xmax
        outer_xmin = -outer_xmax
        outer_ymin = -outer_ymax
        #get fiducial flux
        return _spacecut(indices, d, outer_xmin, outer_xmax, outer_ymin, outer_ymax, scale_to_m=scale_to_m)

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

def parsecml():
    parser = argparse.ArgumentParser(description="Verify genev.")
    parser.add_argument("patterns", metavar="M", type=str, nargs='+', help="Input file list.")
    #parser.add_argument("-f", "--flux", metavar="F", dest="fluxpatterns", type=str, nargs='+', help=".")
    parser.add_argument("-p", "--plot", metavar="P", dest="plot", type=str, choices=["flux", "rate", "ratio"], help="Choose which plot to make.")
    args = parser.parse_args()
    return args

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

def sanitize(args):
    for p in args.patterns:
        if len(glob.glob(p)) < 1:
            raise Exception("Input file does not exist or no matches to the pattern.", p)
    return

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

def main():
    #ROOT.gROOT.ProcessLine(".L rootrackerhelper.cxx++")
    args = parsecml()
    sanitize(args)
    if args.plot == "flux":
        plot_flux(args)
    elif args.plot == "rate":
        plot_interaction_rate(args)
    elif args.plot == "ratio":
        plot_flux_int_xsec(args)
    return

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

if __name__ == "__main__":
    main()

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


import numpy as np

'''
Lower body model
	state vector
	x = [
		root: 			3D(6)	p(x,y,z), u(x,y,z), 	 	,

		l_root: 		3D(6)						theta(x,y,z), 	w(x,y,z),
		l-hip: 			3D(6) 						theta(x,y,z), 	w(x,y,z)	,
		l-knee: 		1D(2) 						theta(z), 		w(z)		,
		l-ankle: 		2D(4) 						theta(y,z), 	w(y,z)		,

        r_root:         3D(6)                       theta(x,y,z),   w(x,y,z),
		r-hip: 			3D(6)						theta(x,y,z), 	w(x,y,z)	,
		r-knee: 		1D(2)						theta(z), 		w(z)		,
		r-ankle: 		2D(4) 						theta(y,z), 	w(y,z)		,

		D_l_root_hip: 	1D(1)	L												,
		D_l_hip_knee: 	1D(1)	L												,
		D_l_knee_ankle: 	1D(1)	L												,
		D_l_ankle_foot: 	1D(1)	L												,
		D_r_root_hip: 	1D(1)	L												,
		D_r_hip_knee: 	1D(1)	L												,
		D_r_knee_ankle: 	1D(1)	L												,
		D_r_ankle_foot: 	1D(1)	L												,
	]
    cal: 42 + 8
	total vector length(DP): 50

	measurement vector
	y = [
		x_root: 		3D(3)
		x_l_hip: 		3D(3)
		x_l-knee:		3D(3)
		x_l-ankle:		3D(3)
		x_l-foot:		3D(3)
		x_r-hip:		3D(3)
		x_r-knee:		3D(3)
		x_r-ankle:		3D(3)
		x_r-foot:		3D(3)
	]
	total vector length(MP): 27
'''
class ukf_Lower_Params:
    def __init__(self, init_mean, init_cov=[1e-6, 1e-4, 1e-6, 1e-4, 1e-6, 1e-1, 1e-6, 1e-1, 1e-6, 1e-4, 1e-9, 1e-4, 100]):
    	self.set_trans_covariance(init_cov[:11])
    	self.set_obs_covariance(init_cov[11])
    	self.set_mean(init_mean)
    	self.set_init_trans_cov(init_cov[12])
    	self.set_trans_matrix()

    def set_state_covariance(self, init_cov):
        root_p = init_cov[0]
        root_v = init_cov[1]
        root_l_t = root_r_t = init_cov[2]
        root_l_w = root_r_w = init_cov[3]
        hip_l_t = hip_r_t = init_cov[4]
        hip_l_w = hip_r_w = init_cov[5]
        knee_l_t = knee_r_t = init_cov[6]
        knee_l_w = knee_r_w = init_cov[7]
        ankle_l_t = ankle_r_t = init_cov[8]
        ankle_l_w = ankle_r_w = init_cov[9]
        link_length = init_cov[10]
        self.state_cov_list = [root_p,root_v,root_l_t,root_l_w,hip_l_t,hip_l_w,knee_l_t,knee_l_w,ankle_l_t,ankle_l_w,root_r_t,root_r_w,hip_r_t,hip_r_w,knee_r_t,knee_r_w,ankle_r_t,ankle_r_w,link_length]

    def set_state_dim(self):
        root_dim = 3
        root_l_dim =  root_r_dim = 3
        hip_l_dim = hip_r_dim = 3
        knee_l_dim = knee_r_dim = 1
        ankle_l_dim = ankle_r_dim = 2
        link_length = 8
        self.state_cov_dim_list = [root_dim,root_dim,root_l_dim,root_l_dim,hip_l_dim,hip_l_dim,knee_l_dim,knee_l_dim,ankle_l_dim,ankle_l_dim,root_r_dim,root_r_dim,hip_r_dim,hip_r_dim,knee_r_dim,knee_r_dim,ankle_r_dim,ankle_r_dim,link_length]
        self.state_cov_total_dim = sum(self.state_cov_dim_list)

    def gen_trans_covariance(self):
        tmp = np.eye(self.state_cov_total_dim)

        idx = 0
        for i in range(len(self.state_cov_list)):
            for j in range(self.state_cov_dim_list[i]):
                tmp[idx+j][idx+j] = self.state_cov_list[i]
            idx = idx + self.state_cov_dim_list[i]

        self.trans_cov = tmp

    def set_trans_covariance(self, init_cov):
        self.set_state_covariance(init_cov)
        self.set_state_dim()
        self.gen_trans_covariance()

    def set_obs_covariance(self, init_obs_cov_factor):
        self.obs_cov_dim = 27
        self.obs_cov_factor = init_obs_cov_factor
        self.obs_cov = np.eye(self.obs_cov_dim)*self.obs_cov_factor

    def set_mean(self, init_mean):
        self.mean = init_mean

    def set_init_trans_cov(self, t_factor):
    	self.init_trans_cov_factor = t_factor
    	self.init_trans_cov = self.trans_cov * self.init_trans_cov_factor

    def set_trans_matrix(self):
        self.fps = 10
        self.is_velocity = [0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0]
        tmp = np.eye(self.state_cov_total_dim)

        idx = 0
        for i in range(len(self.state_cov_list)):
        	if self.is_velocity[i] == 1:
	        	for j in range(self.state_cov_dim_list[i]):
	        		tmp[idx+j-self.state_cov_dim_list[i]][idx+j] = 1/self.fps
        	idx = idx + self.state_cov_dim_list[i]

        self.trans_matrx = tmp

'''
Upper body model
    state vector
    x = [
        root:               3D(6)   p(x,y,z), u(x,y,z),                             ;
        s_root:             3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        spine_naval:        3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_spine_chest:      3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_shoulder:         3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_shoulder_center:  3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_elbow:            3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_wrist_u:          3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_hand:             3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_wrist_d:          3D(6)                       theta(x,y,z),   w(x,y,z),   ;

        r_spine_chest:      3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_shoulder:         3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_shoulder_center:  3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_elbow:            3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_wrist_u:          3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_hand:             3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_wrist_d:          3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        u_spine_chest:      3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        neck:               3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        head:               3D(6)                       theta(x,y,z),   w(x,y,z),   ;

        l_nose:             3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        l_eye:              3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_nose:             3D(6)                       theta(x,y,z),   w(x,y,z),   ;
        r_eye:              3D(6)                       theta(x,y,z),   w(x,y,z),   ;

        D_root_spine:       1D(1)                                                   ,
        D_spine_chest:      1D(1)                                                   ,
        D_l_chest_sh:       1D(1)                                                   ,
        D_l_sh_shc:         1D(1)                                                   ,
        D_l_shc_elbow:      1D(1)                                                   ,
        D_l_elbow_wrist:    1D(1)                                                   ,
        D_l_wrist_hand      1D(1)                                                   ,
        D_l_hand_handtip:   1D(1)                                                   ,
        D_l_wrist_thumb:    1D(1)                                                   ,
        D_r_chest_sh:       1D(1)                                                   ,

        D_r_sh_shc:         1D(1)                                                   ,
        D_r_shc_elbow:      1D(1)                                                   ,
        D_r_elbow_wrist:    1D(1)                                                   ,
        D_r_wrist_hand:     1D(1)                                                   ,
        D_r_hand_handtip:   1D(1)                                                   ,
        D_r_wrist_thumb:    1D(1)                                                   ,
        D_chest_neck:       1D(1)                                                   ,
        D_neck_head:        1D(1)                                                   ,
        D_head_nose:        1D(1)                                                   ,
        D_l_nose_eye:       1D(1)                                                   ,

        D_l_eye_ear:        1D(1)                                                   ,
        D_r_nose_eye:       1D(1)                                                   ,
        D_r_eye_ear:        1D(1)                                                   ,
    ]
    cal: 24*6 - 8 + 23 = 144 + 23 = 167
    total vector length(DP): 167

    measurement vector
    y = [
        root:               3D(3)       0
        spine_naval:        3D(3)       1
        spine_chest:        3D(3)       2
        l_shoulder:         3D(3)       4
        l_shoulder_center:  3D(3)       5
        l_elbow:            3D(3)       6
        l_wrist:            3D(3)       7
        l_hand:             3D(3)       8
        l_handtip:          3D(3)       9
        l_thumb:            3D(3)       10

        r_shoulder:         3D(3)       11
        r_shoulder_center:  3D(3)       12
        r_elbow:            3D(3)       13
        r_wrist:            3D(3)       14
        r_hand:             3D(3)       15
        r_handtip:          3D(3)       16
        r_thumb:            3D(3)       17
        neck:               3D(3)       3
        head:               3D(3)       26
        nose:               3D(3)       27

        l_eye:              3D(3)       28
        l_ear:              3D(3)       29
        r_eye:              3D(3)       30
        r_ear:              3D(3)       31
    ]
    cal: 24 * 3 = 72
    total vector length(MP): 72
'''

class ukf_Upper_Params:
    def __init__(self, init_mean, init_cov=[]):
        self.set_trans_covariance(init_cov[:35])
        self.set_obs_covariance(init_cov[35])
        self.set_mean(init_mean)
        self.set_init_trans_cov(init_cov[36])
        self.set_trans_matrix()

    def set_state_covariance(self, init_cov):
        root_p,                 root_v              = init_cov[0:2]
        s_root_t,               s_root_w            = init_cov[2:4]
        spine_naval_t,          spine_naval_w       = init_cov[4:6]

        l_spine_chest_t,        l_spine_chest_w     = init_cov[6:8]
        l_shoulder_t,           l_shoulder_w        = init_cov[8:10]
        l_shoulder_center_t,    l_shoulder_center_w = init_cov[10:12]
        l_elbow_t,              l_elbow_w           = init_cov[12:14]
        l_wrist_u_t,            l_wrist_u_w         = init_cov[14:16]
        l_hand_t,               l_hand_w            = init_cov[16:18]
        l_wrist_d_t,            l_wrist_d_w         = init_cov[18:20]

        r_spine_chest_t,        r_spine_chest_w     = init_cov[6:8]
        r_shoulder_t,           r_shoulder_w        = init_cov[8:10]
        r_shoulder_center_t,    r_shoulder_center_w = init_cov[10:12]
        r_elbow_t,              r_elbow_w           = init_cov[12:14]
        r_wrist_u_t,            r_wrist_u_w         = init_cov[14:16]
        r_hand_t,               r_hand_w            = init_cov[16:18]
        r_wrist_d_t,            r_wrist_d_w         = init_cov[18:20]

        u_spine_chest_t,        u_spine_chest_w     = init_cov[20:22]
        neck_t,                 neck_w              = init_cov[22:24]
        head_t,                 head_w              = init_cov[24:26]
        l_nose_t,               l_nose_w            = init_cov[26:28]
        l_eye_t,                l_eye_w             = init_cov[28:30]
        r_nose_t,               r_nose_w            = init_cov[30:32]
        r_eye_t,                r_eye_w             = init_cov[32:34]

        link_length                                 = init_cov[34]

        self.state_cov_list = [root_p,root_v,s_root_t,s_root_w,spine_naval_t,spine_naval_w,l_spine_chest_t,l_spine_chest_w,l_shoulder_t,l_shoulder_w,l_shoulder_center_t,l_shoulder_center_w,l_elbow_t,l_elbow_w,l_wrist_u_t,l_wrist_u_w,l_hand_t,l_hand_w,l_wrist_d_t,l_wrist_d_w,r_spine_chest_t,r_spine_chest_w,r_shoulder_t,r_shoulder_w,r_shoulder_center_t,r_shoulder_center_w,r_elbow_t,r_elbow_w,r_wrist_u_t,r_wrist_u_w,r_hand_t,r_hand_w,r_wrist_d_t,r_wrist_d_w,u_spine_chest_t,u_spine_chest_w,neck_t,neck_w,head_t,head_w,l_nose_t,l_nose_w,l_eye_t,l_eye_w,r_nose_t,r_nose_w,r_eye_t,r_eye_w,link_length]

    def set_state_dim(self):
        root_dim = 3
        s_root_dim = 3
        spine_naval_dim = 3

        l_spine_chest_dim = 3
        l_shoulder_dim = 3
        l_shoulder_center_dim = 3
        l_elbow_dim = 3
        l_wrist_u_dim = 3
        l_hand_dim = 3
        l_wrist_d_dim = 3

        r_spine_chest_dim = 3
        r_shoulder_dim = 3
        r_shoulder_center_dim = 3
        r_elbow_dim = 3
        r_wrist_u_dim = 3
        r_hand_dim = 3
        r_wrist_d_dim = 3

        u_spine_chest_dim = 3
        neck_dim = 3
        head_dim = 3
        l_nose_dim = 3
        l_eye_dim = 3
        r_nose_dim = 3
        r_eye_dim = 3

        link_length_dim = 23

        self.state_cov_dim_list = [root_dim,root_dim,s_root_dim,s_root_dim,spine_naval_dim,spine_naval_dim,l_spine_chest_dim,l_spine_chest_dim,l_shoulder_dim,l_shoulder_dim,l_shoulder_center_dim,l_shoulder_center_dim,l_elbow_dim,l_elbow_dim,l_wrist_u_dim,l_wrist_u_dim,l_hand_dim,l_hand_dim,l_wrist_d_dim,l_wrist_d_dim,r_spine_chest_dim,r_spine_chest_dim,r_shoulder_dim,r_shoulder_dim,r_shoulder_center_dim,r_shoulder_center_dim,r_elbow_dim,r_elbow_dim,r_wrist_u_dim,r_wrist_u_dim,r_hand_dim,r_hand_dim,r_wrist_d_dim,r_wrist_d_dim,u_spine_chest_dim,u_spine_chest_dim,neck_dim,neck_dim,head_dim,head_dim,l_nose_dim,l_nose_dim,l_eye_dim,l_eye_dim,r_nose_dim,r_nose_dim,r_eye_dim,r_eye_dim,link_length_dim]
        self.state_cov_total_dim = sum(self.state_cov_dim_list)

    def gen_trans_covariance(self):
        tmp = np.eye(self.state_cov_total_dim)

        idx = 0
        for i in range(len(self.state_cov_list)):
            for j in range(self.state_cov_dim_list[i]):
                tmp[idx+j][idx+j] = self.state_cov_list[i]
            idx = idx + self.state_cov_dim_list[i]

        self.trans_cov = tmp

    def set_trans_covariance(self, init_cov):
        self.set_state_covariance(init_cov)
        self.set_state_dim()
        self.gen_trans_covariance()

    def set_obs_covariance(self, init_obs_cov_factor):
        self.obs_cov_dim = 72
        self.obs_cov_factor = init_obs_cov_factor
        self.obs_cov = np.eye(self.obs_cov_dim)*self.obs_cov_factor

    def set_mean(self, init_mean):
        self.mean = init_mean

    def set_init_trans_cov(self, t_factor):
        self.init_trans_cov_factor = t_factor
        self.init_trans_cov = self.trans_cov * self.init_trans_cov_factor

    def set_trans_matrix(self):
        self.fps = 10
        self.is_velocity = [0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,0]
        tmp = np.eye(self.state_cov_total_dim)

        idx = 0
        for i in range(len(self.state_cov_list)):
            if self.is_velocity[i] == 1:
                for j in range(self.state_cov_dim_list[i]):
                    tmp[idx+j-self.state_cov_dim_list[i]][idx+j] = 1/self.fps
            idx = idx + self.state_cov_dim_list[i]

        self.trans_matrx = tmp