#!/usr/bin/env python

""" A set of tools for modelling the antenna reponse and for callibrating the amplitude vs frequency of the antennas
based on pyCRtools. see Schellart et al. Detecting cosmic rays with the LOFAR radio telescope,  and Nelles et al. Calibrating the absolute amplitude scale for air showers measured at LOFAR

Note: LBA_ant_calibrator still needs some work.

author: Brian hare
"""

##internal
import glob
from pickle import load
import datetime
from os.path import  dirname, abspath

##external
import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
from scipy.interpolate import RegularGridInterpolator, pchip_interpolate, PchipInterpolator
from scipy.io import loadmat

##mine
from LoLIM.utilities import processed_data_dir, MetaData_directory, RTD, SId_to_Sname, v_air
import LoLIM.transmogrify as tmf
import LoLIM.IO.metadata as md

# def get_absolute_calibration(frequencies, freq_mode='30-90'):
#     """this function returns the absolute calibration given a set of frequencies (in Hz). Note that this is the INVERSE of the factor in pyCRtools. You should multiply this by the jones matrix to set the absolute correction
#     Note this is spefically derived for 80m cable lengths, which is consistant with relative calibraiton below.

#     If frequency is outside 30-80 MHz. A is assumed to be constant outside 30-80. Cable lengths and RCU gain are correct between 10-90 MHz. Assumed to be constant outside 10-90.

#     Returns 0 where frequencies<0

#     freq-mode: can be '30-90' or '10-90'. This sets which RCU setting is used. This should probably be the same as the data set being used, but probably doesn't matter.
#           The calibration was derived with 30-80, which is default."""

#     A = [

#         ]

class aartfaac_LBA_model:

    def __init__(self, model_loc=None, R=700, C=15e-12, mode="LBA_OUTER"):
        if model_loc is None:
            model_loc = dirname(abspath(__file__)) + '/AARTFAAC_LBA_MODEL/'

        self.model_loc = model_loc
        self.antenna_mode = mode

        self.fine_model_loc = self.model_loc + "LBA_core_fine"
        self.course_model_loc = self.model_loc + "LBA_core"

        ## get base data, assume is same between models
        data_fine = loadmat( self.fine_model_loc , variable_names=['Positions', 'Theta', 'Phi', 'Freq', 'Zant', 'g_z'] ) ## Zant is antenna impedences, g_z is angular dependence
        self.AART_ant_positions = data_fine['Positions']
        self.AART_thetas = data_fine['Theta'][0]  ## zenithal angle in degrees
        self.AART_Phis = data_fine['Phi'][0]      ## azimuthal angle in degrees

        self.AART_fineFreqs = data_fine['Freq'][0]
        fine_gz = data_fine['g_z']
        fine_Zant =  data_fine['Zant']

        data_course = loadmat( self.course_model_loc , variable_names=['Freq', 'Zant', 'g_z'] )
        self.AART_courseFreqs = data_course['Freq'][0]
        course_gz = data_course['g_z']
        course_Zant =  data_course['Zant']

        ### figure out how to fold fine and course models together
        fineFreq_low = self.AART_fineFreqs[0]
        # fineFreq_high = self.AART_fineFreqs[-1]

        freqs = []
        course_range_1 = [0]
        fine_range = []
        course_range_2 = []

        ## first we search AART_courseFreqs to find highest frequency lower than fine-range, adding them to the frequency list
        for i in range(len(self.AART_courseFreqs)):
            Fc = self.AART_courseFreqs[i]
            if Fc < fineFreq_low:
                freqs.append(Fc)
            else:
                break
        course_range_1.append(i)
        fine_range.append(i)

        ## add all fine frequencies
        freqs += list(self.AART_fineFreqs)
        fine_range.append(len(freqs))
        course_range_2.append(len(freqs))

        ## find first course frequency greater than fine frequencies
        for i in range(i, len(self.AART_courseFreqs)):
            if self.AART_courseFreqs[i] > freqs[-1]:
                break

        ## add rest of course frequencies to list
        freqs += list(self.AART_courseFreqs[i:])
        course_range_2.append(len(freqs))

        self.all_frequencies = np.array(freqs, dtype=np.double)
        self.course_frequencyRange_low = np.array(course_range_1)
        self.fine_frequencyRange = np.array(fine_range)
        self.course_frequencyRange_high = np.array(course_range_2)

        #### combine antenna impedences
        num_ants = len(fine_Zant)
        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)
        num_freqs = len(self.all_frequencies)

        self.antenna_Z = np.empty([num_ants, num_ants, num_freqs ], dtype=np.complex)

        self.antenna_Z[:,:, :self.course_frequencyRange_low[1]] = course_Zant[:,:,  :self.course_frequencyRange_low[1]]
        self.antenna_Z[:,:, self.fine_frequencyRange[0]:self.fine_frequencyRange[1]] = fine_Zant[:,:,  :]
        self.antenna_Z[:,:, self.fine_frequencyRange[1]:] = course_Zant[:,:, self.fine_frequencyRange[1]-len(self.all_frequencies):]

        self.total_impedence = np.empty((num_ants, num_ants, num_freqs), dtype=np.complex)
        self.set_RC( R, C ) ## this sets self.total_impedence

        ### now we combine total G_z, which is voltage on antenna
        self.total_Gz = np.empty( (num_ants,2,num_thetas,num_phis,num_freqs), dtype=np.complex )

        # for ant_i in range(num_ants):
            # for pol_i in range(2):
        self.total_Gz[:,:, :,:,                             : self.course_frequencyRange_low[1]] = course_gz[:,:, :,:,    : self.course_frequencyRange_low[1]]
        self.total_Gz[:,:, :,:, self.fine_frequencyRange[0] : self.fine_frequencyRange[1]]       = fine_gz  [:,:, :,:,    : ]
        self.total_Gz[:,:, :,:, self.fine_frequencyRange[1] :]                                   = course_gz[:,:, :,:, self.fine_frequencyRange[1]-len(self.all_frequencies): ]

    def get_LNA_filter(self):
        """given the antenna mode, return two things. First is np.array of indecies of LNAs that are on, second is np.array of LNAs that are off"""

        station_ant_positions = []
        for station in ['CS002', 'CS003', 'CS004', 'CS005', 'CS006', 'CS007']:
            station_ant_positions.append(  md.convertITRFToLocal( md.getItrfAntennaPosition(station, self.antenna_mode) )[::2, :] )

        found_indeces = []
        unfound_indeces = []

        AARFAACC_X_locations = self.AART_ant_positions[0]
        AARFAACC_Y_locations = self.AART_ant_positions[1]

        err2 = 0.1*0.1
        for AAR_i in range(len(AARFAACC_X_locations)):

            found = False
            for stat_ant_XYZ in station_ant_positions:
                for XYZ in stat_ant_XYZ:
                    KX, KY, KZ = XYZ
                    AX = AARFAACC_X_locations[AAR_i]
                    AY = AARFAACC_Y_locations[AAR_i]
                    R2 = (AX-KX)**2 + (AY-KY)**2
                    if R2<err2:
                        found = True
                        break
                if found:
                    break

            if found:
                found_indeces.append(AAR_i)
            else:
                unfound_indeces.append(AAR_i)

        return np.array(found_indeces), np.array(unfound_indeces)

    def set_RC(self, R, C ):
        self.R = R
        self.C = C

        num_ants = len(self.antenna_Z)
        num_freqs = len(self.all_frequencies)

        # functioning_indeces, nonfunctioning_indeces = self.get_LNA_filter()
        # LNA_filter = np.zeros(num_ants, dtype=np.int)
        # LNA_filter[ functioning_indeces ] = 1
        # LNA_Off_filter = np.zeros(num_ants, dtype=np.int)
        # LNA_Off_filter[ nonfunctioning_indeces ] = 1

        tmp_mat = np.empty((num_ants, num_ants), dtype=np.complex)
        ZLMA_matrix =  np.zeros((num_ants, num_ants), dtype=np.complex)

        for Fi in range(num_freqs):
            # set LNA matrix
            LNA = R/(1+self.all_frequencies[Fi]*(2*np.pi*1j*R*C)) # LNA_filter*(R/(1+self.all_frequencies[Fi]*(2*np.pi*1j*R*C)))  +  LNA_Off_filter*( 1/(self.all_frequencies[Fi]*2*np.pi*1j*C) )
            np.fill_diagonal(ZLMA_matrix,  LNA)

            ## caclulate denominator bit
            tmp_mat[:,:] = self.antenna_Z[:,:,Fi]
            tmp_mat += ZLMA_matrix

            ## try sign difference
            # tmp_mat *= -1
            # np.fill_diagonal(tmp_mat,  np.diagonal(tmp_mat)*-1)

            ## invert and multiple
            inv = np.linalg.inv( tmp_mat )
            np.matmul(ZLMA_matrix, inv, out=tmp_mat)

            # finally
            self.total_impedence[:,:,Fi] = tmp_mat

    def loc_to_anti(self, antenna_XY):
        """given an antenna XY location, return internal index of X-dipole of that antenna"""

        KX, KY = antenna_XY

        AARFAACC_X_locations = self.AART_ant_positions[0, 0::2]
        AARFAACC_Y_locations = self.AART_ant_positions[1, 0::2]

        found_index = None
        err2 = 0.1*0.1
        for AAR_i in range(len(AARFAACC_X_locations)):
            AX = AARFAACC_X_locations[AAR_i]
            AY = AARFAACC_Y_locations[AAR_i]
            R2 = (AX-KX)**2 + (AY-KY)**2
            if R2<err2:
                found_index = AAR_i*2
                break
        if found_index is None:
            print("ERROR! AARTFAAC model cannot find your antenna loc:", antenna_XY)
            quit()

        X_ant_i = found_index
        return X_ant_i

    def RCUname_to_anti(self, antenna_name):
        """given RCUid (and known mode), return internal index of the X-dipole of that antenna"""

        station = SId_to_Sname[ int(antenna_name[:3]) ]
        known_positions = md.convertITRFToLocal( md.getItrfAntennaPosition(station, self.antenna_mode) )
        this_antenna_location = known_positions[ int(antenna_name[-3:]) ]

        return self.loc_to_anti( this_antenna_location[0:2] )

    class single_LBA_model:
        def __init__(self, jones_functions, freq_bounds):
            self.jones_functions = jones_functions
            self.freq_bounds = freq_bounds

        def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
            """ if frequencies is numpy array in Hz, zenith and azimuth in degrees, than return numpy array of jones matrices,
            that when doted with [zenithal,azimuthal] component of incidident E-field, then will give [X,Y] voltages on dipoles"""

            return_matrices = np.empty( (len(frequencies), 2,2), dtype=np.complex )

            if zenith<0:
                zenith = 0
            elif zenith > 90:
                zenith = 90

            while azimuth<0:
                azimuth += 360
            while azimuth>360:
                azimuth -= 360

            # good_frequency_filter = np.logical_and(frequencies>=self.freq_bounds[0], frequencies<self.freq_bounds[-1])

            # points = np.empty( (np.sum(good_frequency_filter), 3), dtype=np.double )

            J00,J01 = self.jones_functions[0]
            J10,J11 = self.jones_functions[1]

            for fi,f in enumerate(frequencies):
                if f<=self.freq_bounds[0] or f>=self.freq_bounds[-1]:
                    return_matrices[fi, 0,0] = freq_fill
                    return_matrices[fi, 1,1] = freq_fill
                    return_matrices[fi, 0,1] = 0.0
                    return_matrices[fi, 1,0] = 0.0
                else:
                    return_matrices[fi, 0,0] = J00( (zenith,azimuth,f) )
                    return_matrices[fi, 0,1] = J01( (zenith,azimuth,f) )
                    return_matrices[fi, 1,0] = J10( (zenith,azimuth,f) )
                    return_matrices[fi, 1,1] = J11( (zenith,azimuth,f) )

                    # print(fi)
                    # M = return_matrices[fi]
                    # print("J:", M)
                    # print( M[ 0,0]*M[ 1,1] - M[  0,1]*M[  1,0] )
                    # print()

            return return_matrices

    def get_antenna_model(self, ant_i):
        """note: ant_i should be internal index!"""

        X_ant_i = int(ant_i/2)*2

        KX, KY = self.AART_ant_positions[:, X_ant_i]

        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)
        num_frequencies = len( self.all_frequencies)
        num_antennas = len( self.total_Gz )

    #### some setup
        ## phase shifter to remove geometric delay
        shifter = np.empty([num_thetas, num_phis, num_frequencies ], dtype=np.complex)
        for zi in range( num_thetas ):
            for ai in range( num_phis ):
                Zenith = self.AART_thetas[zi]/RTD
                Azimuth = self.AART_Phis[ai]/RTD
                dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                np.exp( self.all_frequencies*(1j*2*np.pi*(-dt)), out=shifter[zi,ai] )

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = self.all_frequencies[0]
        maxF = self.all_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        ## memory
        tmp1 = np.empty( (num_antennas,num_frequencies), dtype=np.complex )
        tmp2 = np.empty(num_frequencies, dtype=np.complex )
        tmp3 = np.empty(num_frequencies, dtype=np.double )

        def make_interpolant(antenna, polarization):
            """ antenna is 0 or 1 for X or Y, pol is 0 or 1 for zenith or azimithal component"""
            nonlocal tmp1, tmp2, tmp3

            antenna = antenna + X_ant_i

            grid = np.empty([num_thetas, num_phis, len(interpolant_frequencies) ], dtype=np.complex)

            for theta_i in range(num_thetas):
                for phi_i in range(num_phis):

                    ## dot product between voltages and impedences
                    tmp1[:,:]  = self.total_Gz[:, polarization,theta_i,phi_i, :]
                    tmp1 *= self.total_impedence[antenna, :, :] # this shouldn't matter??
                    # tmp1 *= self.total_impedence[:, antenna, :]

                    np.sum( tmp1, axis=0, out = tmp2 )

                    ## shift phase due to arival direction
                    tmp2 *= shifter[theta_i,phi_i]

                    ## interpolate amplitude and phase
                    np.abs(tmp2, out=tmp3)
                    interp_ampltude = pchip_interpolate(self.all_frequencies,tmp3, interpolant_frequencies)

                    # if theta_i==0 and phi_i==0:
                        # print('amp')
                        # plt.plot( interpolant_frequencies, interp_ampltude, 'o' )
                        # plt.plot( self.all_frequencies, tmp3, 'o' )
                        # plt.show()

                    angles = np.angle(tmp2)
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(self.all_frequencies,angles, interpolant_frequencies)

                    # if theta_i==0 and phi_i==0:
                    #     print('angle')
                    #     plt.plot( interpolant_frequencies, interp_angle, 'o' )
                    #     plt.plot( self.all_frequencies, angles, 'o' )
                    #     plt.show()

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out=grid[theta_i,phi_i] )
                    grid[theta_i,phi_i] *= interp_ampltude

            ## correct for different in definition
            if polarization==0: ## zenith points in oppisite directions in two models??
                grid *= -1
            ## and final angle-frequency interpolant
            interpolant = RegularGridInterpolator((self.AART_thetas, self.AART_Phis, interpolant_frequencies),  grid,  bounds_error=False,fill_value=0.0)
            return interpolant

        J00_interpolant = make_interpolant(0, 0)
        J01_interpolant = make_interpolant(0, 1)
        J10_interpolant = make_interpolant(1, 0)
        J11_interpolant = make_interpolant(1, 1)

        return self.single_LBA_model([[J00_interpolant,J01_interpolant],[J10_interpolant,J11_interpolant]],  [interpolant_frequencies[0],interpolant_frequencies[-1]])

    def get_average_antenna_model(self):
        """return model averaged over antenna set"""

        if self.antenna_mode == "LBA_OUTER":
            start_i = 576
            end_i = 576*2
        elif self.antenna_mode == "LBA_INNER":
            start_i = 0
            end_i = 576
        else:
            print("unknown mode:", self.antenna_mode )

        total_num_antennas = len( self.total_Gz )
        num_frequencies = len(self.all_frequencies)
        num_zeniths = len(self.AART_thetas)
        num_azimuths = len(self.AART_Phis)

        ### define a utility function for extracting grid data
        temp_matrix = np.empty( (num_zeniths, num_azimuths, num_frequencies), dtype=np.complex )
        shifterTMP = np.empty( [ num_frequencies ], dtype=np.complex)
        tmp1 = np.empty( (total_num_antennas,num_frequencies), dtype=np.complex )
        def get_ant_i(ant_i, pol_i):
            """given internal antenna index, and polarization, fill temp_matrix with correct response"""
            nonlocal  shifterTMP, tmp1, temp_matrix

            KX, KY = self.AART_ant_positions[:, ant_i]
            for zi, ze in enumerate(self.AART_thetas):
                for ai, az in enumerate(self.AART_Phis):

                    ## dot product between voltages and impedences
                    tmp1[:,:]  = self.total_Gz[:, pol_i,zi,ai, :]
                    tmp1 *= self.total_impedence[ant_i, :, :] # this shouldn't matter??
                    # tmp1 *= self.total_impedence[:, antenna, :]

                    np.sum( tmp1, axis=0, out = temp_matrix[zi,ai, : ] )

                    ## caculate phase shifts
                    Zenith = ze/RTD
                    Azimuth = az/RTD
                    dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                    shifterTMP[:] = self.all_frequencies
                    shifterTMP *= -1j*2*np.pi*dt
                    np.exp( shifterTMP, out=shifterTMP )

                    ## apply the shits
                    temp_matrix[zi,ai, : ] *= shifterTMP

            if pol_i==0: ## zenith points in oppisite directions in two models??
                temp_matrix *= -1

        ### now define a utility function for calculating the average, and interpolating

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = self.all_frequencies[0]
        maxF = self.all_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        temp_AVE_matrix = np.empty( (num_zeniths, num_azimuths, num_frequencies), dtype=np.complex )
        tmp3 = np.empty(num_frequencies, dtype=np.double )
        def calc_func(antenna_pol, field_pol):
            """antenna_pol should be 0 or 1 for X or Y antenna, and field_pol is 0 or 1 for zenithal or azimuthal field"""
            nonlocal temp_AVE_matrix, tmp3

            ## first average
            Npairs = int( (end_i-start_i)/2 )
            for pair_i in range(Npairs):
                ant_i = 2*pair_i + antenna_pol + start_i
                get_ant_i(ant_i, field_pol) ## this fills temp_matrix
                temp_AVE_matrix += temp_matrix ## error prone, but should be okay

                if not np.all( np.isfinite(temp_AVE_matrix) ):
                    print('YO problem!', pair_i, np.all( np.isfinite(temp_matrix) ) )
                    quit()

            temp_AVE_matrix /= Npairs

            upsample_grid = np.empty([num_zeniths, num_azimuths, len(interpolant_frequencies) ], dtype=np.complex)

            ## now interpolate frequencies
            for zi, ze in enumerate(self.AART_thetas):
                for ai, az in enumerate(self.AART_Phis):

                    np.abs(temp_AVE_matrix[zi,ai, : ], out=tmp3)
                    interp_ampltude = pchip_interpolate(self.all_frequencies, tmp3, interpolant_frequencies)

                    angles = np.angle(temp_AVE_matrix[zi,ai, : ])
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(self.all_frequencies, angles, interpolant_frequencies)

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out = upsample_grid[zi,ai, :] )
                    upsample_grid[zi,ai, :] *= interp_ampltude

            ## and final angle-frequency interpolant
            return RegularGridInterpolator((self.AART_thetas, self.AART_Phis, interpolant_frequencies),  upsample_grid,  bounds_error=False, fill_value=0.0)

        ### now we actually make the responses
        J00 = calc_func(0, 0)
        J01 = calc_func(0, 1)
        J10 = calc_func(1, 0)
        J11 = calc_func(1, 1)
        return self.single_LBA_model([[J00,J01],[J10,J11]],  [interpolant_frequencies[0],interpolant_frequencies[-1]])

    def get_antenna_locs(self):
        """return two arrays. First is X locations of all antennas. Second is Y locations. Each location is in pairs, the first is the X dipole, second is Y dipole"""
        return self.AART_ant_positions[0], self.AART_ant_positions[1]

    def get_grid_info(self):
        """return info about the internal grid. returns three real-valued numpy arrays: frequencies [Hz], zenithal [degrees], and azimuthal [degrees]"""
        return self.all_frequencies, self.AART_thetas, self.AART_Phis

    def get_response_grid(self, antenna_i, frequency_i, ant_pol_i, field_pol_i):
        """given an antenna_i, and frequency_i, return response of ant_pol_i (0 for X 1 for Y dipole) to field_pol_i (0 is zenithal, 1 is azimuthal). Response is a 2D matrix, first index is zenith angle, second is azimuthal"""

        total_antenna_i = int(antenna_i/2)*2 + ant_pol_i

        #### some setup
        KX, KY = self.AART_ant_positions[:, total_antenna_i]

        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)

        frequency = self.all_frequencies[ frequency_i ]

        ret = np.empty([num_thetas, num_phis], dtype=np.complex)

        ## the calculation
        for zenith_i in range( num_thetas ):
            for azimuth_i in range( num_phis ):

                ## phase shifter to remove geometric delay
                Zenith = self.AART_thetas[zenith_i]/RTD
                Azimuth = self.AART_Phis[azimuth_i]/RTD
                dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                phase_shift = np.exp( frequency*(1j*2*np.pi*(-dt)) )

                ## the antenna response
                response = np.dot( self.total_impedence[total_antenna_i, :, frequency_i] , self.total_Gz[:, field_pol_i, zenith_i, azimuth_i, frequency_i] )

                ret[ zenith_i,  azimuth_i] = response*phase_shift

        if field_pol_i==0: ## zenith points in oppisite directions in two models??
            ret *= -1

        return ret

class aartfaac_average_LBA_model:

    def __init__(self, mode="LBA_OUTER"):
        if mode == "LBA_OUTER":
            fname = MetaData_directory+"/lofar/antenna_response_model/AARTFAAC_AVE_LBAOUTER_R700_C17.npz"
        else:
            print('mode unknown:', mode)
            quit()

        data = np.load(fname)
        initial_frequencies = data['freqs']
        zeniths = data['zeniths']
        azimuths = data['azimuths']

        original_J00 = data['J00']
        original_J01 = data['J01']
        original_J10 = data['J10']
        original_J11 = data['J11']

        num_zeniths  = len(zeniths)
        num_azimuths = len(azimuths)

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = initial_frequencies[0]
        maxF = initial_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        # print('arg', initial_frequencies.shape, original_J00.shape)

        tmp = np.empty(len(initial_frequencies), dtype=np.double )
        def upsample_and_interpolate(GRID):
            nonlocal tmp

            upsample_grid = np.empty([num_zeniths, num_azimuths, num_interpolant_freqs ], dtype=np.complex)

            ## now interpolate frequencies
            for zi, ze in enumerate(zeniths):
                for ai, az in enumerate(azimuths):

                    np.abs(GRID[zi,ai, : ], out=tmp)
                    interp_ampltude = pchip_interpolate(initial_frequencies, tmp, interpolant_frequencies)

                    angles = np.angle(GRID[zi,ai, : ])
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(initial_frequencies, angles, interpolant_frequencies)

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out = upsample_grid[zi,ai, :] )
                    upsample_grid[zi,ai, :] *= interp_ampltude

            ## and final angle-frequency interpolant

            return RegularGridInterpolator((zeniths, azimuths, interpolant_frequencies),  upsample_grid,  bounds_error=False, fill_value=0.0)

        self.freq_bounds = [ interpolant_frequencies[0], interpolant_frequencies[-1] ]

        self.J00 = upsample_and_interpolate( original_J00 )
        self.J01 = upsample_and_interpolate( original_J01 )
        self.J10 = upsample_and_interpolate( original_J10 )
        self.J11 = upsample_and_interpolate( original_J11 )

    def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
        """ if frequencies is numpy array in Hz, zenith and azimuth in degrees, than return numpy array of jones matrices,
        that when doted with [zenithal,azimuthal] component of incidident E-field, then will give [X,Y] voltages on dipoles"""

        return_matrices = np.empty( (len(frequencies), 2,2), dtype=np.complex )

        if zenith<0:
            zenith = 0
        elif zenith > 90:
            zenith = 90

        while azimuth<0:
            azimuth += 360
        while azimuth>360:
            azimuth -= 360

        # good_frequency_filter = np.logical_and(frequencies>=self.freq_bounds[0], frequencies<self.freq_bounds[-1])

        # points = np.empty( (np.sum(good_frequency_filter), 3), dtype=np.double )

        for fi,f in enumerate(frequencies):
            if f<self.freq_bounds[0] or f>self.freq_bounds[-1]:
                return_matrices[fi, 0,0] = freq_fill
                return_matrices[fi, 1,1] = freq_fill
                return_matrices[fi, 0,1] = 0.0
                return_matrices[fi, 1,0] = 0.0
            else:
                return_matrices[fi, 0,0] = self.J00( (zenith,azimuth,f) )
                return_matrices[fi, 0,1] = self.J01( (zenith,azimuth,f) )
                return_matrices[fi, 1,0] = self.J10( (zenith,azimuth,f) )
                return_matrices[fi, 1,1] = self.J11( (zenith,azimuth,f) )

                # print(fi)
                # M = return_matrices[fi]
                # print("J:", M)
                # print( M[ 0,0]*M[ 1,1] - M[  0,1]*M[  1,0] )
                # print()

        return return_matrices

    def save_to_text(self):

        freq_grid = np.linspace(10,90, num=int((90-10)/1)+1 )
        zenith_grid = np.linspace(0,90, num=int(90/5)+1 )
        azimuth_grid = np.linspace(0,360, num=int(360/10)+1 )

        J00_out = open('./J00_text.txt', 'w')
        J01_out = open('./J01_text.txt', 'w')
        J10_out = open('./J10_text.txt', 'w')
        J11_out = open('./J11_text.txt', 'w')

        J00_out.write('voltage on X-dipole to zenithal field. Azimuthal=0 points 45 degrees south from East.\n')
        J00_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J01_out.write('voltage on X-dipole to azimuthal field. Azimuthal=0 points 45 degrees south from East.\n')
        J01_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J10_out.write('voltage on Y-dipole to zenithal field\. Azimuthal=0 points 45 degrees south from East.n')
        J10_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J11_out.write('voltage on Y-dipole to azimuthal field. Azimuthal=0 points 45 degrees south from East.\n')
        J11_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        for f in freq_grid:
            for z in zenith_grid:
                for a in azimuth_grid:
                    jonesy = self.Jones_Matrices( [f*1e6],z,a+45 )
                    start = str(f)+ ' ' + str(z)+' '+str(a)+' '
                    J00_out.write(start+str(np.real(jonesy[0,0,0])) + ' ' + str(np.imag(jonesy[0,0,0])) +'\n')
                    J01_out.write(start+str(np.real(jonesy[0,0,1])) + ' ' + str(np.imag(jonesy[0,0,1])) +'\n')
                    J10_out.write(start+str(np.real(jonesy[0,1,0])) + ' ' + str(np.imag(jonesy[0,1,0])) +'\n')
                    J11_out.write(start+str(np.real(jonesy[0,1,1])) + ' ' + str(np.imag(jonesy[0,1,1])) +'\n')

class calibrated_AARTFAAC_model:
    """returns the AARTFAAC model multiplied by katies cal."""

    def __init__(self):
        self.AARTFAAC = aartfaac_average_LBA_model()

        calibration = [42484.88872879, 41519.47733373, 39694.16854372, 38435.36963118,
                       36974.13596039, 35819.34454985, 35072.53901876, 33960.74197721,
                       32944.65142405, 32112.44688046, 31590.52174516, 30433.52586868,
                       29756.92041985, 28880.95581826, 28204.37711364, 27646.83132496,
                       27058.11219529, 26693.43614747, 25952.89684011, 25517.72346825,
                       25220.5602153 , 24782.62398452, 24349.88724441, 23755.78027908,
                       23289.45439997, 22800.30039329, 22176.17569116, 21619.22598549,
                       21341.84386379, 21368.96227069, 22032.38713345, 22127.96304476,
                       22678.25184787, 24670.14420208, 25703.00236409, 24792.9736945,
                       23817.8752368 , 22912.1048524 , 22393.11267467, 21844.38535201,
                       20831.10515113, 20220.27315072, 19223.64518551, 18099.23669312,
                       17569.76228552, 16298.99230863, 14845.6897604 , 13488.79405579,
                       11966.53709451, 10895.66530218,  9703.34502309]

        cal_frequencies = np.arange(30e6, 80.5e6, 1e6)
        self.calibration_interpolator = PchipInterpolator(cal_frequencies, calibration, extrapolate=True  )

    def get_calibrator(self, frequencies):
        return self.calibration_interpolator( frequencies )

    def Jones_ONLY(self, frequencies, zenith, azimuth, freq_fill=1.0):
        return self.AARTFAAC.Jones_Matrices( frequencies, zenith, azimuth, freq_fill )

    def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
        """return calibrated jones matrices"""
        JM = self.AARTFAAC.Jones_Matrices( frequencies, zenith, azimuth, freq_fill )
        C = self.get_calibrator( frequencies )
        JM[:, 0,0] *= C
        JM[:, 0,1] *= C
        JM[:, 1,0] *= C
        JM[:, 1,1] *= C
        return JM

def invert_2X2_matrix_list( matrices ):
    """ if matrices is an array of 2x2 matrices, then return the array of inverse matrices """
    num = len(matrices)
    out = np.zeros( (num, 2,2), dtype=matrices.dtype)

    out[:, 0,0] = matrices[:, 1,1]
    out[:, 0,1] = -matrices[:, 0,1]
    out[:, 1,0] = -matrices[:, 1,0]
    out[:, 1,1] = matrices[:, 0,0]

    determinants = matrices[:, 0,0]*matrices[:, 1,1] - matrices[:, 0,1]*matrices[:, 1,0]

    out /= determinants[:, np.newaxis, np.newaxis]

    return out

def fourier_series( x, p):
    """Evaluates a partial Fourier series

        F(x) \\approx \\frac{a_{0}}{2} + \\sum_{n=1}^{\\mathrm{order}} a_{n} \\sin(nx) + b_{n} \\cos(nx)
    """

    r = p[0] / 2

    order = int( (len(p) - 1) / 2 )

    for i in range(order):

        n = i + 1

        r += p[2*i + 1] * np.sin(n * x) + p[2*i + 2] * np.cos(n * x)

    return r

def getGalaxyCalibrationData(antenna_noise_power, timestamp, antenna_type="outer"):
    """return factor to correct for amplitude shifts. Essenturally returns sqrt( P_{expected} / P_{measured} ). Where P is noise power.
    for antenna_type outer it returns factor for Y/X dipoles, for "inner" returns "X/Y". antenna_noise_power is an array of measured powers for each antenna.
    Even/odd indecies should be Y/X dipole for outer and oppisite for inner.
    timestamp should be posix timestamp"""

    longitude = 6.869837540/RTD

    ## this is in outer order:  Y,X
    # coefficients_lba = [ np.array([ 0.01489468, -0.00129305,  0.00089477, -0.00020722, -0.00046507]),   ## for Y antennas
    #                       np.array( [ 0.01347391, -0.00088765,  0.00059822,  0.00011678, -0.00039787] )  ] ## for X antennas
    coefficients_lba = [ np.array( [ 3.85712631e+01, -2.17182149e+00,  1.68114451e+00 , 4.24076969e-01,  -9.24289199e-01,  2.11372242e-01, 1.09281884e-01, -1.74674795e-01, 3.70793388e-03] ),   ## for Y antennas
                         np.array( [38.13314007, -3.02861767 , 2.11558435, -0.30123627, -0.94641864, -0.14297615,  0.05037442, -0.04133833, -0.11443689])  ] ## for X antennas

    # Convert timestamp to datetime object
    t = datetime.datetime.utcfromtimestamp(timestamp)
    # Calculate JD(UT1)
    ut = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.))
    # Calculate JD(TT)
    dtt = tmf.delta_tt_utc(tmf.date2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.)))
    tt = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + (float(t.second) + dtt / 3600.)) / 24.))
    # Calculate Local Apparant Sidereal Time
    last = tmf.rad2circle(tmf.last(ut, tt, longitude))

    galactic_noise_power =[  fourier_series(last, coefficients_lba[0]),  fourier_series(last, coefficients_lba[1])  ]

    antenna_noise_power[ antenna_noise_power==0 ] = np.nan

    if antenna_type == 'outer':
        Y_measured_powers = antenna_noise_power[0::2]
        X_measured_powers = antenna_noise_power[1::2]

        Y_expected_power = galactic_noise_power[0]
        X_expected_power = galactic_noise_power[1]

    else:
        X_measured_powers = antenna_noise_power[0::2]
        Y_measured_powers = antenna_noise_power[1::2]

        X_expected_power = galactic_noise_power[0]
        Y_expected_power = galactic_noise_power[1]

    ## note this should make new arrays
    Y_factors = Y_expected_power / Y_measured_powers
    X_factors = X_expected_power / X_measured_powers

    np.sqrt(Y_factors, out=Y_factors)
    np.sqrt(X_factors, out=X_factors)

    if antenna_type == 'outer':
        return Y_factors, X_factors
    else:
        return X_factors, Y_factors


import numpy as np
import torch
from PIL import Image

from argparse import ArgumentParser
import base64

from model import *

DEFAULT_OUTPUT = 'output.jpg'
DEFAULT_MODEL = 'model.pt'

def get_opts():
    parser = ArgumentParser()
    parser.add_argument('hash', help='Base64-encoded hash', type=str)
    parser.add_argument('--model', help='Model checkpoint', type=str, default=DEFAULT_MODEL)
    parser.add_argument('--output', help='Output filename', type=str, default=DEFAULT_OUTPUT)
    return parser.parse_args()

def main():
    opts = get_opts()

    # load model
    model = Model()
    # for a single image, faster to do this on CPU
    model.load_state_dict(torch.load(opts.model, map_location=torch.device('cpu')))
    model.eval()

    hash_tensor = torch.tensor(np.array(list(base64.b64decode(opts.hash)), dtype=np.uint8))
    with torch.no_grad():
        # batch size 1
        inverted = model(hash_tensor.unsqueeze(0))[0]
        # convert from CHW to HWC and to uint8
        inverted = np.clip(inverted.permute(1, 2, 0).numpy(), 0, 255).astype(np.uint8)
    Image.fromarray(inverted).save(opts.output)

if __name__ == '__main__':
    main()


#!/usr/bin/env python3
"""Convert multiple pairs."""

import warnings
from pathlib import Path
from functools import partial
from multiprocessing import Pool, cpu_count

import yaml
import torch
import numpy as np
import soundfile as sf
from argparse import ArgumentParser
from tqdm import tqdm

from data import load_wav, log_mel_spectrogram, plot_mel, plot_attn
from data.feature_extract import FeatureExtractor
from models import load_pretrained_wav2vec

def parse_args():
    """Parse command-line arguments."""
    parser = ArgumentParser()
    parser.add_argument("info_path", type=str)
    parser.add_argument("output_dir", type=str, default=".")
    parser.add_argument("-c", "--ckpt_path",
                        default="checkpoints/cpc-cpc.pt")
    parser.add_argument("-s", "--src_feat_name", default="cpc")
    parser.add_argument("-r", "--ref_feat_name", default="cpc")
    parser.add_argument("-w", "--wav2vec_path", default=None)
    parser.add_argument("-v", "--vocoder_path", default=None)

    parser.add_argument("--sample_rate", type=int, default=16000)

    return vars(parser.parse_args())

def main(
    info_path,
    output_dir,
    ckpt_path,
    src_feat_name,
    ref_feat_name,
    wav2vec_path,
    vocoder_path,
    sample_rate,
    **kwargs,
):
    """Main function."""

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    src_feat_model = FeatureExtractor(src_feat_name, wav2vec_path, device)

    ref_feat_model = FeatureExtractor(ref_feat_name, wav2vec_path, device)

    print(f"[INFO] {src_feat_name} is loaded")

    model = torch.jit.load(ckpt_path).to(device).eval()
    print("[INFO] VC Model is loaded from", ckpt_path)

    vocoder = torch.jit.load(vocoder_path).to(device).eval()
    print("[INFO] Vocoder is loaded from", vocoder_path)

    path2wav = partial(load_wav, sample_rate=sample_rate, trim=True)

    with open(info_path) as f:
        infos = yaml.load(f, Loader=yaml.FullLoader)

    out_mels = []
    attns = []
    with Pool(cpu_count()) as pool:
        for pair_name, pair in tqdm(infos.items()):
            src_wav = load_wav(pair["source"], sample_rate, trim=True)
            src_wav = torch.FloatTensor(src_wav).to(device)

            tgt_wavs = pool.map(path2wav, pair["target"])
            tgt_wavs = [torch.FloatTensor(tgt_wav).to(device)
                        for tgt_wav in tgt_wavs]

            with torch.no_grad():
                tgt_mels = ref_feat_model.get_feature(tgt_wavs)
                src_mel = (ref_feat_model.get_feature([src_wav])[0].transpose(
                        0, 1).unsqueeze(0).to(device))
                tgt_mels = [tgt_mel.cpu() for tgt_mel in tgt_mels]
                tgt_mel = np.concatenate(tgt_mels, axis=0)
                tgt_mel = torch.FloatTensor(tgt_mel.T).unsqueeze(0).to(device)
                src_feat = src_feat_model.get_feature([src_wav])[
                    0].unsqueeze(0)
                out_mel, attn = model(src_feat, tgt_mel)

                out_mel = out_mel.transpose(1, 2).squeeze(0)
                out_mels.append(out_mel)
                attns.append(attn)

        # print(f"[INFO] Pair {pair_name} converted")
    # out_mel: batch_size, time_stamp, mel_dim
    del model
    del src_feat_model
    del ref_feat_model
    print("[INFO] Generating waveforms...")
    batch_size = 10
    total = len(out_mels)
    out_wavs = []
    pbar = tqdm(total=len(out_mels), ncols=0, unit="wavs")
    with torch.no_grad():
        for i in range(0, total, batch_size):
            out_wavs.extend(vocoder.generate(out_mels[i:i+batch_size]))
            pbar.update(min(batch_size, total))
        if total % batch_size != 0 and total > batch_size:
            out_wavs.extend(vocoder.generate(
                out_mels[total - total % batch_size:]))
            pbar.update(total % batch_size)
    pbar.close()

    print("[INFO] Waveforms generated")

    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    for pair_name, out_mel, out_wav, attn in tqdm(zip(
        infos.keys(), out_mels, out_wavs, attns
    )):
        out_wav = out_wav.cpu().numpy()
        out_path = Path(out_dir, pair_name)

        plot_mel(out_mel, filename=out_path.with_suffix(".mel.png"))
        plot_attn(attn, filename=out_path.with_suffix(".attn.png"))
        sf.write(out_path.with_suffix(".wav"), out_wav, sample_rate)

if __name__ == "__main__":
    warnings.filterwarnings("ignore")
    main(**parse_args())


# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training helper functions that are shared across tasks."""

import contextlib
import functools
import operator
import signal
import typing
from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union

from absl import logging
import dataclasses
import flax
import gin
import jax
import jax.numpy as jnp
import numpy as np
import optax

from gfsa import jax_util
from gfsa.datasets import data_loading

@jax_util.register_dataclass_pytree
@dataclasses.dataclass
class ExampleWithMetadata:
  """Stores an example or batch of examples.

  Attributes:
    epoch: Integer representing the epoch that this example comes from.
    example_id: Integer ID uniquely identifying this example in the dataset.
    example: The example itself.
    mask: Array that is True for actual examples, False for padding.
    static_metadata: Metadata about this example or batch that should result in
      a new `jit` XLA computation (i.e. padded shapes).
  """
  epoch: Any
  example_id: Any
  example: Any
  mask: jax_util.NDArray = np.array(True)
  static_metadata: Any = None

@jax_util.register_dataclass_pytree
@dataclasses.dataclass
class RatioMetric:
  """A ratio, where numerator and denominator should be summed separately.

  Attributes:
    numerator: Numerator of the metric.
    denominator: Denominator of the metric.
  """
  numerator: jax_util.NDArray
  denominator: jax_util.NDArray

MetricValue = Union[float, jax_util.NDArray, RatioMetric]

# A loss function is a callable (model, example, static_metadata)
#                                 -> (loss, metrics)
# pyformat: disable
LossFunWithMetrics = Callable[
    [Any, Any, Any],
    Tuple[jax_util.NDArray, Dict[str, MetricValue]]]
# pyformat: enable

# A validation function is a callable (replicated_model) -> (objective, metrics)
# where model is a tree of ShardedDeviceArrays, and objective is the value we
# want to make decrease.
ValidationFunction = Callable[[Any], Tuple[float, Dict[str, MetricValue]]]

def device_broadcast(x, num_devices):
  """Broadcast a value to all devices."""
  return jax.pmap(lambda _: x)(jnp.arange(num_devices))

def _parallel_train_step(
    optimizer,
    batched_examples,
    static_batch_metadata,
    loss_fn,
    max_global_norm = None,
    **optimizer_hyper_params,
):
  """Train the model for one step in parallel across devices.

  Args:
    optimizer: Optimizer that tracks the model and parameter state. Should be
      replicated to each device, i.e. should contain ShardedDeviceArrays with a
      leading axis (num_devices, ...) but with the same content on each device.
    batched_examples: A structure of NDArrays representing a batch of examples.
      Should have two leading batch dimensions: (num_devices,
        batch_size_per_device, ...)
    static_batch_metadata: Metadata about this batch, which will be shared
      across all batched examples. Each value of this results in a separate
      XLA-compiled module.
    loss_fn: Task-specific non-batched loss function to apply. Should take the
      current model (optimizer.target) and an example from batched_examples, and
      return a tuple of the current loss (as a scalar) and a dictionary from
      string names to metric values (also scalars, or RatioMetrics).
    max_global_norm: Maximum global norm to clip gradients to. Should be a
      scalar, which will be broadcast automatically.
    **optimizer_hyper_params: Hyperparameters to pass to the optimizer's
      `apply_gradient` function, which will be broadcast across devices
      automatically.

  Returns:
    Tuple (updated_optimizer, grads_ok, metrics). Metrics will be as returned by
    loss_fn, with an extra elements "loss". All metrics will be averaged
    across all elements of the batch. Both optimizer and metrics will contain
    ShardedDeviceArrays that are identical across devices. grads_ok will be
    a replicated bool ndarray that is True if the gradients were finite.
  """

  def batched_loss_fn(model):
    """Apply loss function across a batch of examples."""
    loss, metrics = jax.vmap(loss_fn, (None, 0, None))(model, batched_examples,
                                                       static_batch_metadata)
    return jnp.mean(loss), metrics

  # Compute gradients of loss, along with metrics.
  (loss, metrics), grads = jax.value_and_grad(
      batched_loss_fn, has_aux=True)(
          optimizer.target)
  metrics["loss"] = loss

  # Exchange average gradients and metrics across devices.
  agg_grads = jax.lax.pmean(grads, "devices")
  agg_metrics = {}
  for k, v in metrics.items():
    if isinstance(v, RatioMetric):
      num = jax.lax.psum(jnp.sum(v.numerator), "devices")
      denom = jax.lax.psum(jnp.sum(v.denominator), "devices")
      new_value = num / denom
    else:
      # Use nanmean to aggregate bare floats.
      new_value = jnp.nanmean(jax.lax.all_gather(v, "devices"))
    agg_metrics[k] = new_value

  # Compute global norm and possibly clip.
  global_norm = optax.global_norm(agg_grads)
  agg_metrics["gradient_global_norm"] = global_norm
  if max_global_norm is not None:
    should_clip = global_norm > max_global_norm
    agg_grads = jax.tree_map(
        lambda g: jnp.where(should_clip, g * max_global_norm / global_norm, g),
        agg_grads)
    agg_metrics["gradient_was_clipped"] = should_clip.astype("float32")

  # Check for non-finite gradients.
  grads_ok = jnp.all(
      jnp.stack([jnp.all(jnp.isfinite(x)) for x in jax.tree_leaves(agg_grads)]))

  # Apply updates.
  updated_optimizer = optimizer.apply_gradient(agg_grads,
                                               **optimizer_hyper_params)

  return updated_optimizer, grads_ok, agg_metrics, agg_grads

def _build_parallel_train_step():
  """Builds an accelerated version of the train step function."""
  # We need to wrap and unwrap so that the final function can be called with
  # keyword arguments, but we still maintain the proper axes.

  @functools.partial(
      jax.pmap,
      axis_name="devices",
      in_axes=(0, 0, None, None, None, None),
      static_broadcasted_argnums=(2, 3))
  def wrapped(optimizer, batched_examples, static_batch_metadata, loss_fn,
              max_global_norm, optimizer_hyper_params):
    return _parallel_train_step(optimizer, batched_examples,
                                static_batch_metadata, loss_fn, max_global_norm,
                                **optimizer_hyper_params)

  @functools.wraps(_parallel_train_step)
  def wrapper(optimizer, batched_examples, static_batch_metadata, loss_fn,
              max_global_norm, **optimizer_hyper_params):
    return wrapped(optimizer, batched_examples, static_batch_metadata, loss_fn,
                   max_global_norm, optimizer_hyper_params)

  return wrapper

# The primary version of the training step, with the associated jit cache.
parallel_train_step = _build_parallel_train_step()

def warmup_train_step(
    optimizer,
    batched_example,
    static_batch_metadata,
    loss_fn,
    optimizer_is_replicated = False,
    profile = False,
    runner=None,
):
  """Run a fake train step to warm up JIT cache.

  Args:
    optimizer: Optimizer that tracks the model and parameter state.
    batched_example: A structure of NDArrays representing a batch of examples.
    static_batch_metadata: Metadata about the batch, which will be shared across
      all batched examples.
    loss_fn: Task-specific non-batched loss function to apply. Should take the
      current model (optimizer.target) and an example from batched_examples, and
      return a tuple of the current loss (as a scalar) and a dictionary from
      string names to metric values (also scalars).
    optimizer_is_replicated: Whether optimizer is already replicated.
    profile: Whether to enable profiling during warmup.
    runner: If profile=True, the runner to use when profiling.
  """
  num_devices = jax.local_device_count()
  if optimizer_is_replicated:
    replicated_optimizer = optimizer
  else:
    replicated_optimizer = device_broadcast(optimizer, num_devices)

  (replicated_optimizer,
   batched_example) = jax.tree_map(jax.device_put,
                                   (replicated_optimizer, batched_example))

  try:
    max_global_norm = gin.query_parameter(
        "train_util.training_loop.max_global_norm")
  except ValueError:
    max_global_norm = None

  def go():
    # Note that value for learning_rate is arbitrary, but we pass it here to
    # warm up the jit cache (since we are passing a learning rate at training
    # time).
    res = parallel_train_step(
        replicated_optimizer,
        batched_example,
        static_batch_metadata,
        loss_fn,
        max_global_norm=max_global_norm,
        learning_rate=0.0)
    jax.tree_map(lambda x: x.block_until_ready(), res)

  if profile:
    stats = runner.try_run_and_profile(go, catch_resource_exhausted=False)
    logging.info("Warmed up train step with stats: %s", stats)
  else:
    go()
    logging.info("Warmed up train step")

def build_averaging_validator(
    loss_fn,
    valid_iterator_factory,
    objective_metric_name = None,
    include_total_counts = False,
    prefetch = True,
):
  """Validate by computing averages over a validation set.

  Args:
    loss_fn: Loss function for the task.
    valid_iterator_factory: Constructs iterators of batched examples from the
      validation set, with two batch axes. To iterate over a fixed part of the
      validation set, consider using build_one_pass_iterator_factory. To
      randomly sample from a validation set, you can use something like
      `lambda: itertools.islice(validation_iterator, num_batches)`.
    objective_metric_name: Name of the metric that is the objective value.
    include_total_counts: Whether to report numerator and denominator separately
      for RatioMetric objects, along with the "validation_total_example_count"
      metric.
    prefetch: Whether to prefetch validation examples.

  Returns:
    Validation function that runs loss_fn and aggregates the results, reporting
    the loss as the objective, and using sum to accumulate metrics.
  """
  if objective_metric_name is None:
    objective_metric_name = "loss"

  @functools.partial(
      jax.pmap, axis_name="devices", static_broadcasted_argnums=3)
  def parallel_metrics_batch(model, batched_examples, batch_mask,
                             static_metadata):
    loss, metrics = jax.vmap(loss_fn, (None, 0, None))(model, batched_examples,
                                                       static_metadata)
    metrics["loss"] = loss
    metrics = jax.tree_map(
        lambda x: jnp.where(batch_mask, x, jnp.zeros_like(x)), metrics)
    metrics = jax.tree_map(lambda x: jax.lax.psum(jnp.sum(x), "devices"),
                           metrics)
    return metrics

  def validation_function(model):
    with contextlib.ExitStack() as exit_stack:
      valid_iterator = valid_iterator_factory()
      if prefetch:
        valid_iterator = exit_stack.enter_context(
            data_loading.ThreadedPrefetcher(valid_iterator, 4))
      accumulated = None
      example_count = 0
      for batch in valid_iterator:
        results = parallel_metrics_batch(model, batch.example, batch.mask,
                                         batch.static_metadata)
        metrics = jax.tree_map(float, flax.jax_utils.unreplicate(results))
        metrics["epoch"] = np.sum(batch.epoch)
        if accumulated is None:
          accumulated = metrics
        else:
          accumulated = jax.tree_multimap(operator.add, accumulated, metrics)
        example_count += jnp.count_nonzero(batch.mask)

      assert example_count > 0, "Validation iterator must be nonempty"
      accumulated = typing.cast(Dict[str, Any], accumulated)

      final_metrics = {}
      for k, v in accumulated.items():
        if isinstance(v, RatioMetric):
          final_metrics[k] = v.numerator / v.denominator
          if include_total_counts:
            final_metrics[k + "_numerator"] = v.numerator
            final_metrics[k + "_denominator"] = v.denominator
        else:
          final_metrics[k] = v / example_count

      objective = final_metrics[objective_metric_name]
      if include_total_counts:
        final_metrics["validation_total_example_count"] = example_count
      return (objective, final_metrics)

  return validation_function

@contextlib.contextmanager
def catch_interrupts_once(callback,
                          catch_signals = (signal.SIGINT,
                                                          signal.SIGABRT)):
  # pylint: disable=g-doc-return-or-yield
  """Context manager to catch interrupt signals.

  Only catches the first signal sent, so that repeated interrupts will kill the
  job as normal.

  Args:
    callback: Function to run when the signal is caught the first time.
    catch_signals: Signals to catch.

  Returns:
    A context manager that will catch interrupts inside the block.
  """

  # pylint: enable=g-doc-return-or-yield
  known_signals = {
      signal.SIGINT: "SIGINT",
      signal.SIGABRT: "SIGABRT",
  }

  def _handler(signal_number, frame):
    del frame  # Unused.
    logging.warning("Caught interrupt signal %s",
                    known_signals.get(signal_number, signal_number))
    callback(signal_number)
    _restore_handlers()

  original_handlers = {}
  for signal_number in catch_signals:
    original_handlers[signal_number] = signal.signal(signal_number, _handler)

  already_restored = False

  def _restore_handlers():
    nonlocal already_restored
    if already_restored:
      return
    else:
      already_restored = True
      for signal_number in catch_signals:
        current_handler = signal.signal(signal_number,
                                        original_handlers[signal_number])
        if current_handler is not _handler:
          logging.error(
              "Unexpected active signal handler %s for %s; "
              "expected the signal hander from "
              "`catch_interrupts_once`! Restored to %s anyways.",
              current_handler, known_signals.get(signal_number, signal_number),
              original_handlers[signal_number])

  try:
    yield
  finally:
    _restore_handlers()


# -*-coding:utf-8-*-

import os
import sys

import numpy as np
from PIL import Image
import tensorflow as tf

_MODEL_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'data')

_IMAGE_SIZE = 64
_BATCH_SIZE = 128

_LABEL_MAP = {0: 'drawings', 1: 'hentai',  2: 'neutral', 3: 'porn', 4: 'sexy'}

def standardize(img):
    mean = np.mean(img)
    std = np.std(img)
    img = (img - mean) / std
    return img

def load_image(infilename):
    img = Image.open(infilename)
    img = img.resize((_IMAGE_SIZE, _IMAGE_SIZE))
    img.load()
    data = np.asarray(img, dtype=np.float32)
    data = standardize(data)
    return data

def predict_nsfw(image_path):
    with tf.Session() as sess:
        graph = tf.get_default_graph()
        tf.saved_model.loader.load(
            sess, [tf.saved_model.tag_constants.SERVING], _MODEL_DIR)
        inputs = graph.get_tensor_by_name("input_tensor:0")
        probabilities_op = graph.get_tensor_by_name('softmax_tensor:0')
        class_index_op = graph.get_tensor_by_name('ArgMax:0')
        image_data = load_image(image_path)
        probabilities, class_index = sess.run(
            [probabilities_op, class_index_op],
            feed_dict={inputs: [image_data] * _BATCH_SIZE})
        probabilities_dict = {
            _LABEL_MAP.get(i): l for i, l in enumerate(probabilities[0])}
        pre_label = _LABEL_MAP.get(class_index[0])
        return {"class": pre_label, "probability": probabilities_dict}

if __name__ == '__main__':
    argv = sys.argv
    if(len(argv) < 2):
        print("usage: python rsnsfw.py <image_path>")
    image_path = argv[1]
    print()
    res = predict_nsfw(image_path)
    print(res)


#! /usr/bin/env python

'''
To make the generation of WFSS simulated integrations easier,
combine the 4 stages of the simulator
(seed image generator, disperser, dark prep, obervation generator)
into a single script.

Inputs:
paramfiles - List of yaml filenames. These files should be
             inputs to the simulator. For details on the
             information contained in the yaml files, see
             the readme file associated with the mirage
             github repo:
             https://github.com/spacetelescope/mirage.git

crossing_filter - Name of the crossing filter to be used in
                  conjunction with the grism. All longwave
                  channel filter names are valid entries

module - Name of the NIRCam module to use for the simulation.
         Can be 'A' or 'B'

direction - Dispersion direction. Can be along rows ('R') or
            along columns ('C')

override_dark - If you wish to use a dark current file that
                has already gone through the dark_prep step
                of the pipeline and wish to use that for the
                simulation, set override_dark equal to the
                dark's filename. The dark_prep step will then
                be skipped.

HISTORY:
15 November 2017 - created, Bryan Hilbert
13 July 2018 - updated for name change to Mirage, Bryan Hilbert
'''

import os
import sys
import argparse

from numpy import nanmedian, isfinite
from astropy.io import fits
from NIRCAM_Gsim.grism_seed_disperser import Grism_seed

from .seed_image import catalog_seed_image
from .dark import dark_prep
from .ramp_generator import obs_generator
from .utils import read_fits
from .yaml import yaml_update

nircam_filters = ['F322W2', 'F277W', 'F356W', 'F444W', 'F250M', 'F300M',
                  'F335M', 'F360M', 'F410M', 'F430M', 'F323N', 'F405N',
                  'F466N', 'F470N']

class WFSSSim():
    def __init__(self):
        # Set the MIRAGE_DATA environment variable if it is not
        # set already. This is for users at STScI.
        self.env_var = 'MIRAGE_DATA'
        self.datadir = os.environ.get(self.env_var)
        if self.datadir is None:
            raise ValueError(("WARNING: {} environment variable is not set."
                              "This must be set to the base directory"
                              "containing the darks, cosmic ray, PSF, etc"
                              "input files needed for the simulation."
                              "These files must be downloaded separately"
                              "from the Mirage package.".format(self.env_var)))
        self.paramfiles = None
        self.override_dark = None
        self.crossing_filter = None
        self.module = None
        self.direction = None
        self.prepDark = None
        self.save_dispersed_seed = True
        self.disp_seed_filename = None
        self.extrapolate_SED = False
        self.fullframe_apertures = ["NRCA5_FULL", "NRCB5_FULL"]

    def create(self):
        # Make sure inputs are correct
        self.check_inputs()

        # Loop over the yaml files and create
        # a direct seed image for each
        imseeds = []
        # Create imaging seed images
        for pfile in self.paramfiles:
            cat = catalog_seed_image.Catalog_seed()
            cat.paramfile = pfile
            cat.make_seed()
            imseeds.append(cat.seed_file)

        # Create dispersed seed image from
        # the direct images
        dmode = 'mod{}_{}'.format(self.module,self.direction)
        loc = os.path.join(self.datadir,"nircam/GRISM_NIRCAM/")
        background_file = ("{}_{}_back.fits"
                           .format(self.crossing_filter,dmode))
        disp_seed = Grism_seed(imseeds, self.crossing_filter,
                               dmode, config_path=loc,
                               extrapolate_SED=self.extrapolate_SED)
        disp_seed.observation()
        disp_seed.finalize(Back = background_file)

        # Get gain map
        gainfile = cat.params['Reffiles']['gain']
        gain, gainheader = self.read_gain_file(gainfile)

        # Disperser output is always full frame. Crop to the
        # requested subarray if necessary
        if cat.params['Readout']['array_name'] not in self.fullframe_apertures:
            print("Subarray bounds: {}".format(cat.subarray_bounds))
            print("Dispersed seed image size: {}".format(disp_seed.final.shape))
            disp_seed.final = self.crop_to_subarray(disp_seed.final, cat.subarray_bounds)
            gain = self.crop_to_subarray(gain, cat.subarray_bounds)
            # Segmentation map will be centered in a frame that is larger
            # than full frame by a factor of sqrt(2), so crop appropriately
            segy, segx = cat.seed_segmap.shape
            dx = int((segx - 2048) / 2)
            dy = int((segy - 2048) / 2)
            segbounds = [cat.subarray_bounds[0] + dx, cat.subarray_bounds[1] + dy,
                         cat.subarray_bounds[2] + dx, cat.subarray_bounds[3] + dy]
            cat.seed_segmap = self.crop_to_subarray(cat.seed_segmap, segbounds)

        # Convert seed image to ADU/sec to be consistent
        # with other simulator outputs
        disp_seed.final /= gain

        # Update seed image header to reflect the
        # division by the gain
        cat.seedinfo['units'] = 'ADU/sec'

        # Save the dispersed seed image
        if self.save_dispersed_seed:
            hh00 = fits.PrimaryHDU()
            hh11 = fits.ImageHDU(disp_seed.final)
            hhll = fits.HDUList([hh00,hh11])
            hhll[0].header['units'] = 'ADU/sec'
            if self.disp_seed_filename is None:
                pdir, pf = os.path.split(self.paramfiles[0])
                dname = 'dispersed_seed_image_for_' + pf + '.fits'
                self.disp_seed_filename = os.path.join(pdir, dname)
            hhll.writeto(self.disp_seed_filename, overwrite=True)
            print(("Dispersed seed image saved to {}"
                   .format(self.disp_seed_filename)))

        # Prepare dark current exposure if
        # needed.
        if self.override_dark is None:
            d = dark_prep.DarkPrep()
            d.paramfile = self.paramfiles[0]
            d.prepare()
            obslindark = d.prepDark
        else:
            self.read_dark_product(self.override_dark)
            obslindark = self.prepDark

        # Using the first of the imaging seed image yaml
        # files as a base, adjust to create the yaml file
        # for the creation of the final dispersed
        # integration
        y = yaml_update.YamlUpdate()
        y.file = self.paramfiles[0]
        y.filter = self.crossing_filter
        y.pupil = 'GRISM' + self.direction
        y.outname = ("wfss_dispersed_{}_{}.yaml"
                     .format(dmode,self.crossing_filter))
        y.run()

        # Combine into final observation
        obs = obs_generator.Observation()
        obs.linDark = obslindark
        obs.seed = disp_seed.final
        obs.segmap = cat.seed_segmap
        obs.seedheader = cat.seedinfo
        obs.paramfile = y.outname
        obs.create()

    def read_dark_product(self,file):
        # Read in dark product that was produced
        # by dark_prep.py
        self.prepDark = read_fits.Read_fits()
        self.prepDark.file = file
        self.prepDark.read_astropy()

    def check_inputs(self):
        # Make sure input parameters are good
        if self.module not in ['A','B']:
            self.invalid('module',self.module)
        else:
            self.module = self.module.upper()

        if self.direction not in ['R','C']:
            self.invalid('direction',self.direction)
        else:
            self.direction = self.direction.upper()

        if self.crossing_filter not in nircam_filters:
            self.invalid('crossing_filter',self.crossing_filter)
        else:
            self.crossing_filter = self.crossing_filter.upper()

        if self.override_dark is not None:
            avail = os.path.isfile(self.override_dark)
            if avail == False:
                print(("WARNING: {} does not exist."
                       .format(self.override_dark)))
                sys.exit()

        if len(self.paramfiles) < 2:
            print("WARNING: self.paramfiles must be a list")
            print("of 2 or more yaml files.")
            sys.exit()

    def read_param_file(self, file):
        """
        Read in yaml simulator parameter file

        Parameters:
        -----------
        file -- Name of a yaml file in the proper format
                for mirage

        Returns:
        --------
        Nested dictionary with the yaml file's contents
        """
        import yaml
        try:
            with open(file, 'r') as infile:
                data = yaml.load(infile)
        except:
            raise IOError("WARNING: unable to open {}".format(file))

    def read_gain_file(self, file):
        """
        Read in CRDS-formatted gain reference file

        Paramters:
        ----------
        file -- Name of gain reference file

        Returns:
        --------
        Detector gain map (2d numpy array)
        """
        try:
            with fits.open(file) as h:
                image = h[1].data
                header = h[0].header
        except:
            raise IOError("WARNING: Unable to open gain file: {}".format(file))

        mngain = nanmedian(image)

        # Set pixels with a gain value of 0 equal to mean
        image[image == 0] = mngain
        # Set any pixels with non-finite values equal to mean
        image[~isfinite(image)] = mngain
        return image, header

    def crop_to_subarray(self, data, bounds):
        """
        Crop the given full frame array down to the appropriate
        subarray size and location based on the requested subarray
        name.

        Parameters:
        -----------
        data -- 2d numpy array. Full frame image. (2048 x 2048)
        bounds -- 4-element list containing the full frame indices that
                  define the position of the subarray.
                  [xstart, ystart, xend, yend]

        Returns:
        --------
        Cropped 2d numpy array
        """
        yl, xl = data.shape
        valid = [False, False, False, False]
        valid = [(b>=0 and b<xl) for b in bounds[0:3:2]]
        validy = [(b>=0 and b<yl) for b in bounds[1:4:2]]
        valid.extend(validy)

        print(valid)
        print(bounds)
        print(yl, xl)

        if all(valid):
            return data[bounds[1]:bounds[3] + 1, bounds[0]:bounds[2] + 1]
        else:
            raise ValueError(("WARNING: subarray bounds are outside the "
                              "dimensions of the input array."))

    def read_subarr_defs(self, subfile):
        # read in the file that contains a list of subarray
        # names and positions on the detector
        try:
            subdict = ascii.read(subfile, data_start=1, header_start=0)
            return subdict
        except:
            raise RuntimeError(("Error: could not read in subarray definitions file "
                                "{}".format(subfile)))

    def get_subarr_bounds(self, subname, sdict):
        # find the bounds of the requested subarray
        if subname in sdict['AperName']:
            mtch = subname == sdict['AperName']
            bounds = [sdict['xstart'].data[mtch][0], sdict['ystart'].data[mtch][0],
                      sdict['xend'].data[mtch][0], sdict['yend'].data[mtch][0]]
            return bounds
        else:
            raise ValueError(("WARNING: {} is not a subarray aperture name present "
                              "in the subarray definition file.".format(subname)))

    def invalid(self,field,value):
        print(("WARNING: invalid value for {}: {}"
               .format(field,value)))
        sys.exit()

    def add_options(self,parser = None, usage = None):
        if parser is None:
            parser = argparse.ArgumentParser(usage = usage, description="Wrapper for the creation of WFSS simulated exposures.")
        parser.add_argument("paramfiles",help='List of files describing the input parameters and instrument settings to use. (YAML format).',nargs='+')
        parser.add_argument("--crossing_filter",help = "Name of crossing filter to use in conjunction with the grism.",default=None)
        parser.add_argument("--module",help = "NIRCam module to use for simulation. Use 'A' or 'B'",default=None)
        parser.add_argument("--direction",help = "Direction of dispersion (along rows or along columns). Use 'R' or 'C'",default=None)
        parser.add_argument("--override_dark",help="If supplied, skip the dark preparation step and use the supplied dark to make the exposure", default=None)
        parser.add_argument("--extrapolate_SED", help="If true, the SED created from the filter-averaged magnitudes will be extrapolated to fill the wavelngth range of the grism", action='store_true')
        return parser

if __name__ == '__main__':

    usagestring = 'USAGE: wfss_simualtor.py file1.yaml file2.yaml --crossing_filter F444W --direction R --module A'

    obs = WFSSSim()
    parser = obs.add_options(usage = usagestring)
    args = parser.parse_args(namespace = obs)
    obs.create()


import os
import numpy as np
import torch
from torch.nn.functional import softmax
from tqdm import tqdm

#Imports for Backbone
from genz3d.kpconv.models.architectures import KPFCNN
from genz3d.kpconv.train_SemanticKitti import SemanticKittiConfig

#Imports from 3DGenZ modules
from genz3d.seg.utils.loss import SegmentationLosses
from genz3d.seg.utils.lr_scheduler import LR_Scheduler
from genz3d.seg.utils.metrics import Evaluator
from genz3d.seg.dataloaders import make_data_loader
from genz3d.seg.trainer_class import Trainer_default, main
from genz3d.seg.class_names import  CLASSE_NAMES_SK

class Trainer(Trainer_default):

    def __init__(self, args):
        super().__init__(args)
        # Choose to train on CPU or GPU
        if args.cuda and torch.cuda.is_available():
            self.device = torch.device("cuda:0")
        else:
            self.device = torch.device("cpu")
        # Define Dataloader
        self.N_CLASSES = 19
        self.class_names = CLASSE_NAMES_SK

        #Get the data loader
        self.config = SemanticKittiConfig()
        self.nclass = self.N_CLASSES
        self.train_loader, self.val_loader, training_dataset, test_dataset = make_data_loader(args, config=self.config)
        model = self.get_model(self.config, training_dataset)
        #Load the ConvBasic version
        chkp_path = os.path.join(args.savedir, "chkp_0250.tar")
        print("Load the pre-trained network from {}".format(chkp_path))
        #Filter out the ones that don't match (from classification training)
        checkpoint = torch.load(chkp_path)
        filter_dict = {"head_softmax.mlp.weight":0, "head_softmax.batch_norm.bias":0, "head_softmax_generative.mlp.weight":0, "head_softmax_generative.batch_norm.bias":0}
        pretrained_filtered_dict = {k: v for k, v in checkpoint['model_state_dict'].items() if not k in filter_dict}
        model.load_state_dict(pretrained_filtered_dict, strict=False)
        #Freeze all layers besides the linear layer that gets finetuned
        model.freeze_backbone_fcout()
        model.freeze_bn()

        train_params = [
            {"params": model.get_1x_lr_params(), "lr": args.lr},
            {"params": model.get_10x_lr_params(), "lr": args.lr * 10},
        ]

        # Define Optimizer
        optimizer = torch.optim.SGD(
            train_params,
            momentum=args.momentum,
            weight_decay=args.weight_decay,
            nesterov=args.nesterov,
        )

        # Define Generator and generator optimizer
        embed_feature_size = 0
        self.set_generator(embed_feature_size=embed_feature_size)

        class_weight = torch.ones(self.nclass)
        for l in args.unseen_classes_idx_metric:
            if l-1 >= 0:
                class_weight[l-1] = args.unseen_weight #-1 is necessary for shift between 1 to 20 and 0 to 19
        print("Class weight: {}".format(class_weight))

        if args.cuda:
            class_weight = class_weight.cuda()
        self.criterion = SegmentationLosses(
            weight=class_weight, cuda=args.cuda, bs=self.args.batch_size
        ).build_loss(mode=args.loss_type)
        self.model, self.optimizer = model, optimizer

        #Define Evaluator
        #Shifted seen classes + shifted unseen classes idx:
        shifted_unseen_classes_idx_metric = [l-1 for l in args.unseen_classes_idx_metric]

        ### FOR METRIC COMPUTATION IN ORDER TO GET PERFORMANCES FOR TWO SETS
        shifted_seen_classes_idx_metric = np.arange(20)
        shifted_seen_classes_idx_metric = np.delete(shifted_seen_classes_idx_metric, shifted_unseen_classes_idx_metric).tolist()
        self.evaluator = Evaluator(self.nclass + 1, shifted_seen_classes_idx_metric, shifted_unseen_classes_idx_metric) #Get 1 class more

        # Define lr scheduler
        self.scheduler = LR_Scheduler(
            args.lr_scheduler, args.lr, args.epochs, len(self.train_loader)
        )

        # Using cuda
        if args.cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
            self.model = self.model.cuda()
            self.generator = self.generator.cuda()

        # Resuming checkpoint
        self.best_pred = 0.0
        if args.resume is not None:
            if not os.path.isfile(args.resume):
                raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']

            if args.cuda:
                self.model.module.load_state_dict(checkpoint["state_dict"], strict=False)
            else:
                self.model.load_state_dict(checkpoint["state_dict"])

            self.best_pred = checkpoint['best_pred']
            print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))

        # # Clear start epoch if fine-tuning
        if args.ft:
            args.start_epoch = 0

    def get_model(self, config, training_dataset):
        net = KPFCNN(config, training_dataset.label_values, training_dataset.ignored_labels)
        return net

    def training(self, epoch, args):

        train_loss = 0.0
        self.model.module.train()
        self.model.module.freeze_backbone_fcout()
        self.model.module.freeze_bn()

        tbar = tqdm(self.train_loader)
        num_img_tr = len(self.train_loader)
        i = None
        for i, sample in enumerate(tbar):
            pts, target, embedding = (
                sample.points, sample.labels, sample.attributes
            )

            target = torch.unsqueeze(target, 1)
            if 'cuda' in self.device.type:
                sample.to(self.device)
                target = target.cuda()
                embedding = embedding.cuda()
            self.scheduler(self.optimizer, i, epoch, self.best_pred)

            # ===================real feature extraction=====================
            with torch.no_grad():
                real_features, _ = self.model.module.backbone(
                    batch=sample, config=self.config, return_mid=True)
                target_predicted_backbone = torch.argmax(self.model.module.trained_ss_ll(x=real_features, batch=sample, config=self.config), 1)
            real_features = real_features.permute(1, 0)

            # ===================fake feature generation=====================
            fake_features = torch.zeros(real_features.shape)
            if args.cuda:
                fake_features = fake_features.cuda()
            generator_loss_batch = 0.0
            pi_start = 0
            pi_end = 0

            for _, lenght_i in enumerate(sample.lengths[0]):

                pi_end += lenght_i
                _, real_features_i, target_i, embedding_i, target_predicted_i = pts[pi_start:pi_end],\
                    real_features[:, pi_start:pi_end], target[pi_start:pi_end], embedding[pi_start:pi_end], target_predicted_backbone[pi_start:pi_end]
                generator_loss_sample = 0.0
                ## reduce to real feature size
                real_features_i = (
                    real_features_i.permute(1, 0)
                    .contiguous()
                    .view((-1, args.feature_dim))
                )

                target_i = target_i.view(-1)
                target_predicted_i = target_predicted_i.view(-1)
                fake_features_i = torch.zeros(real_features_i.shape)
                if args.cuda:
                    fake_features_i = fake_features_i.cuda()

                unique_class = torch.unique(target_i)

                ## test if Point Cloud has unseen class pixel, if yes means no training for generator and generated features for the whole PC
                has_unseen_class = False
                for u_class in unique_class:
                    if u_class in args.unseen_classes_idx_metric:
                        has_unseen_class = True

                #Generate for each class the examples
                for idx_in in unique_class:
                    if idx_in != 255:
                        self.optimizer_generator.zero_grad()
                        idx_class = target_i == idx_in
                        real_features_class = real_features_i[idx_class]
                        embedding_class = embedding_i[idx_class]
                        #Noise generation
                        z = torch.rand((embedding_class.shape[0], args.noise_dim))
                        if args.cuda:
                            z = z.cuda()

                        # Generation of the fake Features
                        if args.generator_model == "gmmn":
                            fake_features_class = self.generator(
                                embedding_class, z.float()
                            )
                            fake_feature_train = fake_features_class
                        else:
                            print("Incorrect generator model selection")
                            raise NotImplementedError

                        #Training of the generator, only if in PC is not a single unseen point
                        if (idx_in in args.seen_classes_idx_metric and not has_unseen_class):
                            # in order to avoid CUDA out of memory
                            random_idx = torch.randint(
                                low=0,
                                high=fake_feature_train.shape[0],
                                size=(args.batch_size_generator,),
                            )

                            #Generator loss between fake features und real features
                            g_loss = self.criterion_generator(
                                fake_feature_train[random_idx],
                                real_features_class[random_idx],
                            )
                            generator_loss_sample += g_loss.item()
                            g_loss.backward()
                            self.optimizer_generator.step()

                        fake_features_i[idx_class] = fake_feature_train.clone()
                generator_loss_batch += generator_loss_sample / len(unique_class)

                if args.real_seen_features and not has_unseen_class:
                    #Take the real features if there is no unseen class in PC
                    fake_features[:, pi_start:pi_end] = real_features_i.permute(1, 0)
                else:
                    #Take the generated features if there is an unsenn class in PC
                    fake_features[:, pi_start:pi_end] = fake_features_i.permute(1, 0)

                pi_start = pi_end.detach().clone()

            # Training of the Classifier
            self.optimizer.zero_grad()
            fake_features = fake_features.permute(1, 0).contiguous()
            output = self.model.module.training_generative(x=fake_features.detach())
            target_adapted = self.model.module.labels_adaption(target)
            loss = self.criterion(output.view(-1, self.N_CLASSES), target_adapted.view(-1))
            loss.backward()
            self.optimizer.step()
            train_loss += loss.item()
            # ===================log=====================
            tbar.set_description(
                " G loss:{:.3f}".format(generator_loss_batch)
                + " C loss:{:.3f}".format(train_loss / (i + 1))
            )
            self.writer.add_scalar(
                "train/total_loss_iter", loss.item(), i + num_img_tr * epoch
            )
            self.writer.add_scalar(
                "train/generator_loss", generator_loss_batch, i + num_img_tr * epoch
            )

        self.writer.add_scalar("train/total_loss_epoch", train_loss, epoch)
        print("Pts {}".format(len(pts)))
        print("pts size {}".format(pts[0].size()))
        print("[Epoch: %d, numPCs: %5d]"% (epoch, i * self.args.batch_size + pts[0].size(0)))
        print("Loss: {:.3f}".format(train_loss))

    def validation(self, epoch, args):
        # Important: This validation function only is a rough and random picking method for monitoring purposes during training.
        # A consistent and complete test/valiadtion has to be run with the method provided in the KP-Conv folder.
        self.model.eval()
        self.evaluator.reset()
        tbar = tqdm(self.val_loader, desc="\r")
        test_loss = 0.0
        i = None
        for i, sample in enumerate(tbar):
            pts, target, embedding = (sample.points, sample.labels, sample.attributes)

            target = torch.unsqueeze(target, 1)
            if self.args.cuda:
                sample.to(self.device)
                pts, target, embedding = pts[0].cuda(), target.cuda(), embedding.cuda()

            with torch.no_grad():
                output = self.model.module.forward_generative(batch=sample, config=self.config)
                target_adapted = self.model.module.labels_adaption(target)
                loss = self.criterion(output.view(-1, self.N_CLASSES), target_adapted.view(-1))
                pred = softmax(output, dim=1).data.cpu().numpy()
                bias_matrix = np.zeros(pred.shape)

                for i in args.seen_classes_idx_metric[0:3]:
                    if i >= 0:
                        bias_matrix[:, i-1] = 1.0 * args.bias
                pred = np.argmax(pred - bias_matrix, axis=1)

            test_loss += loss.item()
            tbar.set_description("Test loss: %.3f" % (test_loss / (i + 1)))

            target_adapted = target_adapted.cpu().numpy()
            target = target.cpu().numpy()

            # Add batch sample into evaluator
            self.evaluator.add_batch(np.squeeze(target_adapted), np.squeeze(pred))

        print("IMPORTANT: This validation function is only for MONITORING (see KP-Conv folder for validation function)")
        print("Fast validation (only for monitoring during training):")
        self.validation_metric_logs(test_loss=test_loss, epoch=epoch, num=i * self.args.batch_size + pts.size(0))
        print("IMPORTANT: This validation function is only for MONITORING (see ConvPoint folder for test function)")

if __name__ == "__main__":
    main("sk")


import os
import sys
import re
import json
import numpy as np
import pandas as pd
import argparse
from tqdm import tqdm
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.wcs import WCS
from astropy import log
log.setLevel('ERROR')

this_file_location = os.path.dirname(os.path.abspath(__file__))

parser = argparse.ArgumentParser(
    description=(
        'Generate subsection of the NSA catalogue used by galaxy_utilities. '
        'Requires make_nsa_subsection.py to have been run, and the output '
        'df-nsa.pkl to be present in the same folder as this file. '
        'Also requires a Zooniverse Subject export, and the locations of the '
        'FITS files used in the subject creation process (fits headers are '
        'required for rotations)'
    )
)
default_subjects_csv_location = os.path.join(
    this_file_location,
    'galaxy-builder-subjects.csv'
)
default_montages_location = os.path.join(
    this_file_location,
    '../../gzbuilder_data_prep/montageOutputs'
)

default_subjects_location = os.path.join(
    this_file_location,
    '../../gzbuilder_data_prep/fits_images'
)

parser.add_argument('--subjects', metavar='/path/to/subjects.csv',
                    default=default_subjects_csv_location,
                    type=str, help='Location of Zooniverse Subject export')

parser.add_argument('--montages', metavar='/path/to/montageOutputs',
                    default=default_montages_location, type=str,
                    help='Location of Montage Output folder from subject creation')

parser.add_argument('--images', metavar='/path/to/images',
                    default=default_subjects_location, type=str,
                    help='Location of FITS Images downloaded from SkyServer during subject creation')

args = parser.parse_args()

montage_groups = [
    i for i in os.listdir(args.montages)
    if re.match(r'([0-9.]+)\+([\-0-9.]+)', i) is not None
]

# create a list of coordinates to search
montage_coords = SkyCoord(
    np.array([
        i.groups() for i in (
            re.match(r'([0-9.]+)\+([\-0-9.]+)', i)
            for i in montage_groups
        )
        if i is not None
    ]).astype(float),
    unit=u.degree,
)

subjects = pd.read_csv(
    args.subjects,
).drop_duplicates(subset='subject_id').set_index('subject_id', drop=False)

metadata = subjects.metadata.apply(json.loads).apply(pd.Series)

try:
    df_nsa = pd.read_pickle(os.path.join(this_file_location, 'df_nsa.pkl'))
except FileNotFoundError:
    print('Could not find df_nsa.pkl, please run "make_nsa_subsection.py"')
    sys.exit(0)

def rotation(a):
    return np.array(((np.cos(a), np.sin(a)), (-np.sin(a), np.cos(a))))

def transform_angle(original_frame, montage_frame, angle):
    wcs = dict(original=WCS(original_frame), montaged=WCS(montage_frame))
    angles = dict()
    for k in wcs.keys():
        centre = wcs[k].wcs.crval.copy()
        dec_line = centre + (10/3600, 0)
        # negative as NSA uses (N to E) angles
        line_of_angle = np.dot(
            rotation(np.deg2rad(angle)),
            dec_line - centre
        ) + centre
        centre_pix, angle_line_pix = wcs[k].all_world2pix(
            [centre, line_of_angle],
            0
        )
        new_angle = np.arctan2(*np.flip(angle_line_pix - centre_pix))
        angles[k] = -new_angle
    return angles

def get_input_frames(gal):
    gal_coord = SkyCoord(gal['RA'], gal['DEC'], unit=u.degree)
    idx = np.argmin(montage_coords.separation(gal_coord))
    fpath = os.path.join(args.montages, montage_groups[idx])
    return [
        os.path.join(fpath, i)
        for i in os.listdir(fpath)
        if '.fits' in i
    ]

def get_output_frame(input_frames):
    if len(input_frames) == 1:
        return input_frames[0]
    os.path.join(os.path.dirname(input_frames[0]), 'mosaic.fits')
    return os.path.join(
        os.path.dirname(input_frames[0]),
        'mosaic.fits'
    ).replace('montageGroups', 'montageOutputs')

def get_frames(gal):
    input_frames = get_input_frames(gal)
    output_frame = get_output_frame(input_frames)
    return input_frames, output_frame

if __name__ == '__main__':
    galaxy_info = pd.Series([])
    with tqdm(subjects.index.values, desc='Iterating over subjects') as bar:
        for subject_id in bar:
            # Grab the metadata of the subject we are working on
            subject = subjects.loc[subject_id]
            # And the NSA data for the galaxy (if it's a galaxy with NSA data,
            # otherwise throw an error)
            if metadata.loc[subject_id].get('NSA id', False) is not np.nan:
                try:
                    gal = df_nsa.drop_duplicates(
                        subset='NSAID'
                    ).set_index(
                        'NSAID',
                        drop=False
                    ).loc[
                        int(metadata.loc[subject_id]['NSA id'])
                    ]
                except KeyError:
                    gal = {}
                    raise KeyError(
                        'Metadata does not contain valid NSA id (probably an older galaxy)'
                    )
                original_fits_names, montage_fits_name = get_frames(gal)
                try:
                    angles = transform_angle(
                        original_fits_names[0],
                        montage_fits_name,
                        gal['PETRO_PHI90']
                    )
                except OSError as e:
                    print(subject_id, e)
                    continue
                gal['montage_angle'] = angles['montaged']
                gal['original_angle'] = angles['original']
                # legacy
                gal['angle'] = angles['montaged']
                galaxy_info.loc[subject_id] = gal

    galaxy_info.apply(pd.Series).to_csv('gal-metadata.csv')

import numpy,import astropy,from astropy


import os
#import sys
#import math
import numpy as np
#import subprocess
import time

from config import *

##########     .     ##########

##########     settings     ##########
gamma=2 # 1 =enable, 0 =disable
bright=0.263 # 0 to 1 linear brightness
xs=6 # horizontal scale, integer
ys =xs # vertical scale, integer
bl=8 # bit depth per color, max 8
bright *= xf

startx=100
width=64
starty=100
height=32

#$%&@! 7 MSB - 0 LSB

xl=64 # columns
yl=16 # pjysical rows
zl=2  # screens

##########     prepare color look-up table     ##########
out_lut = np.zeros((zl, 8), dtype="uint32")
for i in range(zl):
    for n in range(8):
        d=0
        if(n&(1<<0)):
            d += out_map[i][0]
        if(n&(1<<1)):
            d += out_map[i][1]
        if(n&(1<<2)):
            d += out_map[i][2]
        out_lut[i, n]=d

#print(out_lut)
print("refresh rate: ", 40*1000*1000/(xf+30)/(yf+3), "hz")

##########     prepare black line + clock     ##########
arr_out_black = np.full((1, xf), clk, dtype="uint32")
front_blank = np.full((13, xf), clk, dtype="uint32")
count=0
timer=9
for x in range(0, xl):
    count+=1
    d2=0

    arr_out_black[0, count] = d2
    count+=1
    arr_out_black[0, count] = d2+clk

##########     spawn arrays     ##########
arr_input= np.zeros((1,32, 64, 3), dtype="uint8")
#           ^ bitplane, y,  x, color
arr_gamma= np.zeros((32, 64, 3), dtype="uint16")
#                   ^ y,  x, color
arr_bitplanes= np.zeros((8, 32, 64, 3), dtype="uint8")
#                    ^ bit,  y,  x, color
arr_sum= np.zeros((8, 32, 64), dtype="uint8")
#              ^ bit,  y,  x
arr_out_data = np.zeros((8, 16, xf), dtype="uint32")
#                   ^ bit,  y, x
d3 = np.zeros((8, 16, xl), dtype="uint32")

#arr_bitmask = np.zeros((8, 32, 64, 3), dtype="uint8") # full size array
#arr_bitmask = np.zeros((8, 1, 64, 3), dtype="uint8") # doesn't hurt, doesn't help
arr_bitmask = np.zeros((8, 1, 1, 1), dtype="uint8") # 2% speed improve
for n in range(0,8):
    arr_bitmask[n, :, :, :] = (1<<n)   #$%&@! 7 MSB - 0 LSB

arr_out_br_clk = np.full((8, 16, xf), clk, dtype="uint32")
#                      ^ bit,  y, x
arr_out_br_clk[:, :, 1:xl*2+1:2] = 0
arr_out_br_clk[:, :, 2:xl*2+2:2] = clk

##########     apply brightness to output data     ##########
def set_brightness(value, array):
    #value *= xf

    out = np.full((8, 16, xf), clk, dtype="uint32") # CLK all high for stability reason
    #                      ^ bit,  y, x
    out[:, :, 1:xl*2+1:2] = 0
    out[:, :, 2:xl*2+2:2] = clk # apply CLK to data section

    for n in range (8):
        if (int(array[n])):
            out[n,:,0:int(bright)] += enable # add brightness bit for FULL lines (don't multplay)
        else:
            out[n,:,0:int(bright*array[n])] += enable # add brightness bit for HALF lines

    # little hack to shift color and brightness one line off
    #out = np.flip(out,0) # flip to put MSB (brightest bitplane) at begining of the array
    out = np.concatenate((out[0, 0, :].ravel(), out.ravel())) # flatten the array and add extra column to move brightness bits one row away from data bits (added later)
    out = np.resize(out, (8, 16, xf)) #, refcheck=False) # resize back to expected size, cut off extra bytes
    out = np.flip(out,0) # and flip back to expected order
    return(out) #$%&@! 7 MSB - 0 LSB

bt_lut = np.array([8, 4, 2, 1, 1/2, 1/4, 1/8, 1/16]) #$%&@! 7 MSB - 0 LSB
#      ^ bitplanes 0, 1, 2, 3,   4,   5,   6,   7

out_multiply = np.array([8, 4, 2, 1, 1, 1, 1, 1]) #$%&@! 7 MSB - 0 LSB

arr_out_br_clk = set_brightness(bright, bt_lut).copy()

##########     apply same brightness to empty line     ##########
arr_out_bright = arr_out_black.copy()
arr_out_bright[0,0:int(bright*bt_lut[7])] += enable

if (0): # old code, that works but looks terrible
    #arr_out_data xf 16 8 to bit7
    arr_out_br_clk[4:8,:,0:int(bright)] += enable # "full" brightness for bitplanes 4 to 7
    arr_out_br_clk[3,0,0:int(bright)] += enable # and first row of 3 bitplane

    arr_out_br_clk[3,1:16,0:int(bright/2)] += enable # half brightness for bitplane 3
    arr_out_br_clk[2,0,0:int(bright/2)] += enable

    arr_out_br_clk[2,1:16,0:int(bright/4)] += enable # quater brightness for bitplane 2
    arr_out_br_clk[1,0,0:int(bright/4)] += enable

    arr_out_br_clk[1,1:16,0:int(bright/8)] += enable # bitplane 1
    arr_out_br_clk[0,0,0:int(bright/8)] += enable

    arr_out_br_clk[0,1:16,0:int(bright/16)] += enable # bitplane 0

while 1:
    time1 = time.time()

    ##########     import image     ##########
    with open('/dev/fb0','rb') as s:
        s.seek(xhost*starty*4, 0)
        raw=s.read(height*ys*xhost*4)
        buffer = np.frombuffer(raw, dtype="uint8").reshape((1, height*ys,xhost,4))

        arr_input = buffer[0, ::ys, startx:startx+xs*width:xs, 0:3]
        s.close()

    ##########     super simple gamma correction     ##########
    if(gamma):
        arr_gamma = arr_input.astype("uint16")
        arr_gamma = arr_gamma**2
        arr_input = np.right_shift(arr_gamma, 8).astype("uint8")

    time2 = time.time()

    ##########     convert image to bitplanes     ##########
    arr_bitplanes = np.bitwise_and(arr_bitmask, arr_input)
    arr_bitplanes[:,:,:,0] = np.where(arr_bitplanes[:,:,:,0], 4, 0)
    arr_bitplanes[:,:,:,1] = np.where(arr_bitplanes[:,:,:,1], 2, 0)
    arr_bitplanes[:,:,:,2] = np.where(arr_bitplanes[:,:,:,2], 1, 0)

    arr_sum = np.bitwise_or(np.bitwise_or(arr_bitplanes[:,:,:,0], arr_bitplanes[:,:,:,1]), arr_bitplanes[:,:,:,2])

    d3 = np.bitwise_or(out_lut[0, arr_sum[:, 0:16, :]], out_lut[1, arr_sum[:, 16:32, :]])

    arr_out_data = arr_out_br_clk.copy()

    arr_out_data[:, :, 1:xl*2+1:2] += d3
    arr_out_data[:, :, 2:xl*2+2:2] += d3

    time3 = time.time()

    ##########     write bitplanes to output framebuffer     ##########
    a7 = arr_out_data[7,:,:]
    a6 = arr_out_data[6,:,:]
    a5 = arr_out_data[5,:,:]
    a4 = arr_out_data[4,:,:]
    a3 = arr_out_data[3,:,:]
    a2 = arr_out_data[2,:,:]
    a1 = arr_out_data[1,:,:]
    a0 = arr_out_data[0,:,:]
    #a = np.concatenate((front_blank, a7, a7, a7, a7, a7, a7, a7, a7, a6, a6, a6, a6, a5, a5, a4, a3, a2, a1, a0, arr_out_bright, arr_out_black))
    # total 318 lines to print out, 360 limit set in bootup config.txt
    #a = np.repeat(np.flip(arr_out_data,0), out_multiply, axis=0).reshape((-1, 240))

    #print(a.shape)
    #print(front_blank.shape)
    #print(arr_out_bright.shape)
    #print(arr_out_black.shape)
    #exit()

    #a = np.concatenate((front_blank, np.repeat(np.flip(arr_out_data,0), out_multiply, axis=0).reshape((-1, 240)), arr_out_bright, arr_out_black))
    a = np.concatenate((front_blank, np.repeat(np.flip(arr_out_data,0), out_multiply, axis=0).reshape((-1, 240)), arr_out_bright, arr_out_black))

    with open('/dev/fb1','wb') as f:
        #a.tofile(f)

        #front_blank.tofile(f)
        a.tofile(f)   #$%&@! 7 MSB - 0 LSB
        #arr_out_bright.tofile(f)
        #arr_out_black.tofile(f)
        f.close()

    ##########     calculate total fps and possible fps if divided into parallel tasks     ##########
    time4 = time.time()

    t1 = "%.3f" % (1.0/(time2-time1))
    t2 = "%.3f" % (1.0/(time3-time2))
    t3 = "%.3f" % (1.0/(time4-time3))
    t4 = "%.3f" % (1.0/(time4-time1))

    #time.sleep(0.05)
    time5 = time.time()
    t5 = "%.3f" % (1.0/(time5-time1))

    if (timer):
        timer-=1
    else:
        timer=40
        print("fps:", t4, "imp", t1, "conv", t2, "send", t3, "real", t5)
    #print("input ", 1.0/(draw_time-start_time), " draw", 1.0/(time.time()-draw_time))