# Copyright (c) 2012-2018, University of Strathclyde
# Authors: Lawrence T. Campbell
# License: BSD-3-Clause
#
# Python script containing the code to write the input files for
# Puffin.
#
# The three files are:
#
#   1) The main input file, describing free parameters and the
#      integration sampling
#
#   2) The beam input file, which can describe multiple electron
#      beams if necessary
#
#   3) The seed input file, which can describe multiple seeds
#
# This approach is (hopefully) quite flexible, allowing Puffin to
# generate a simple, more conventional FEL, or a more esoteric
# configuration with multiple seeds and electron beams, with different
# powers, frequencies, energies, and distributions.
#
# In addition, for more fun, a lattice file may be constructed to
# describe a series of undulator modules seperated by chicane/slippage
# sections and/or quads. The name of this lattice file should be entered
# in the 'lattfile' variable in the input file.
#
#
#     DESCRIBE UNDULATOR LATTICE FILE
#
#
# Also note that if aw is decreased in an undulator module you should
# ensure the radiation mesh is fine enough to model the higher frequency
# content which will arise.
#
# The strength of the dispersive chicanes can be given in the 'Dfact'
# variable in the main input file. Currently, only this one strength
# factor may be specified which will be shared by all chicanes.
#
# An undulator taper can also be added in the main input file.
# Presently, if a lattice file is used, this taper will be applied
# to ALL undulator modules, and the taper will be applied to the
# undulator parameter of each module. Only a linear taper may be
# supplied.
#
# - Lawrence Campbell
#   University of Strathclyde
#   July 2013

# simple setup
  # define beam and undulator - frame is assumed the same as the beam and undulator
  # scale parameters

# - OR -

# advanced setup
  # define frame
  # define beam in frame
    # use same beam as reference beam?
  # scale parameters

import math
import beamClass
import seedClass
import fBoolean
import numpy as np

torf = fBoolean.torf

pow = math.pow
pi = math.pi
sqrt = math.sqrt
log = math.log
bm = beamClass.bm
sd = seedClass.sd

# Physical constants

m_e = 9.109e-31
q_e = 1.602e-19
eps_0 = 8.85e-12
c = 2.997924e8

# File names for input files

inputfile = 'example.in'
beamfile = 'beam_file.in'
seedfile = 'seed_file.in'
lattFile = ''

nbeams = 1
nseeds = 1

### Sampling

elms_per_wave = 20         # Field nodes per resonant wavelength
emps_per_wave = 24         # Electron macroparticles per resonant wavelength
steps_per_per = 30        # should be roughly 4-5* elms_per_wave

# Undulator and beam parameters
# Below similar to CLARA parameters

Er = 240e6            # Beam energy (reference particle)
gamma = Er / ( (m_e * pow(c,2) / q_e) ) # Rel. factor
aw = 1.01               # PEAK wiggler parameter
lambda_w = 0.0275        # Undulator period
N_w = 200              # Number of undulator periods
undtype = 'planepole'
ux = 0
uy = 1

# RMS wiggler param depends on undulator type / polarization:

if (undtype == 'planepole'):
  ux = 0
  uy = 1
  awrms = aw / np.sqrt(2)

elif (undtype == 'curved'):
  ux = 0
  uy = 1
  awrms = aw / np.sqrt(2.)

elif (undtype == 'helical'):
  ux = 1
  uy = 1
  awrms = aw

else:
  awrms = aw / np.sqrt(2.) * np.sqrt(ux**2 + uy**2)

qFlatTopZ2 = np.array([1])         # =1 if flat top current profile, else gaussian.
qHardEdgeX = np.array([0])         # =1 if disk (circle) in transverse plane, else gaussian.
qRoundZ2 = np.array([1])           # If rounding off edges of flat top in z2
sigRound_lam = np.array([6])       # Sigma of gaussian used to round off the flat top edges, in resonant wavelengths
bEnOscMag = np.array([0])          # Magnitude of oscillation on beam energy | Units of gamma
bEnOscFr = np.array([7E3])         # Frequency (2pi/lambda, where lambda in units of ct) of oscillation on beam energy
#E = 300e6            # Beam energy
#gamma = E / (m_e * pow(c,2)) # Rel. factor
sig_gamma = np.array([0.0004])      # Energy spread (relative to Er i.e. sig_gam / gamma_r)

Ej = np.array([240e6])
eratio = Ej / Er # [ij / Er for ij in Ej] # Rel. factor
gammaj = Ej / ( (m_e * pow(c,2) / q_e) ) # [ij / ( (m_e * pow(c,2) / q_e) ) for ij in Ej] # Rel. factor
emitx = np.array([1e-6 / gamma])       # Unnormalised Emittance in x
emity = np.array([1e-6 / gamma])       # Unnormalised Emittance in y
Q = np.array([0.1e-9])                # Charge

k_w = 2 * pi / lambda_w             # Get wiggler wavenumber
sigt = np.array([250e-15]) # 0.005570423008216 / c        # Get sigma in t dimension
sigz = c * sigt                     # Convert sigma in t to z

#if
ff = np.sqrt(2)

k_beta = aw * k_w / ( ff * gamma )  # Betatron wavenumber
N = [ij / q_e for ij in Q]                         # Number of real electrons in pulse
lambda_r = lambda_w / (2 * np.power(gamma,2)) * (1 + np.power(awrms,2))
                                    # ^ Resonant wavelength

sigRoundz = sigRound_lam * lambda_r

k_betax = k_beta
k_betay = k_beta   #...else strong...

alphax = np.array([0.])
alphay = np.array([0.])

#sigx = np.sqrt(emitx / k_betax)     # Beam standard deviation in x...
#sigy = np.sqrt(emity / k_betay)     # ...and y

sigx = np.array([50e-6])
sigy = np.array([50e-6])

sig_av = np.sqrt((np.power(sigx,2) + np.power(sigy,2))/2.0)

########################################
# Beam area and FEL parameter

if qHardEdgeX == 1:
    r_av = sig_av                 # Hard edged circle in x and y
else:
    r_av = np.sqrt(2.) * sig_av  # Gaussian dist in x and y

tArea = np.pi * np.power(r_av,2)          # Tranverse beam area

if qFlatTopZ2[0] == 1:
    if qRoundZ2[0] == 1:
        lArea = np.sqrt(2*pi) * sigRoundz  + sigz     # flat top + gaussian
    else:
        lArea = sigz                  # longitudinal integral over charge dist (flat top)
else:
    lArea = sqrt(2*pi) * sigz     # longitudinal integral over charge dist(gaussian)

n_p = N / (tArea * lArea)              # Electron number density
wp = np.sqrt(np.power(q_e,2) * n_p / (eps_0 * m_e) )         # plasma frequency

rho = 1.0 / gamma * np.power((aw * wp[0] / ( 4.0 * c * k_w )),(2.0/3.0))  # FEL parameter

#######################################
# Scaled parameters for Puffin

lambda_z2 = 4*pi*rho              # Resonant wavelength in z2
Lg = lambda_w / lambda_z2         # Gain length
Lc = lambda_r / lambda_z2         # Cooperation length
zbarprop = N_w * lambda_z2        # Length of undulator in zbar (gain lengths)
sigz2 = sigz / Lc                 # Length of pulse in z2 (cooperation lengths)

sigRoundZ2 = sigRoundz / Lc       # Sigma of tail off in z2 (cooperation lengths)

bEnOscFr = bEnOscFr * Lc

beta = sqrt(np.power(gamma,2) - 1.0 - np.power(awrms,2))/gamma    # Average velocity over c
eta = (1-beta)/beta                                # Scaled average velocity (of reference electron)

k_beta_bar = k_beta * Lg                           # Scaled betatron wavenumber
emitx_bar = emitx / (rho * Lc)                       # Scaled emittance
emity_bar = emity / (rho * Lc)                       # Scaled emittance
Z_R = pi * np.power(r_av,2) / lambda_r                  # Rayleigh range
Z_bar_R = np.power(r_av,2) / (Lg * Lc) / (4.0 * rho)      # Scaled Rayleigh Range
B = np.power((2 * Z_bar_R),(3.0/2.0))                       # Saldin diffraction parameter

NL = N/sigz2 * lambda_z2                           # electrons per radiation period

Anoise = 6 * sqrt(pi) * rho / (NL * sqrt(log(NL/rho)))  # Spontaneous noise estimate (in scaled units)
Acse = 16 * np.power(rho,2)                             # CSE estimate for flat-top current

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

if (undtype == 'planepole'):
  kbxn = 0.
  kbyn = aw / 2. / np.sqrt(2) / rho / gamma # 'natural' focusing wavenumber

elif (undtype == 'curved'):
  kxu = np.sqrt(eta / 8. / rho**2)
  kyu = np.sqrt(eta / 8. / rho**2)
  kbxn = aw / np.sqrt(2.*eta) / gamma *kxu
  kbyn = aw / np.sqrt(2.*eta) / gamma *kyu

elif (undtype == 'helical'):
  kbxn = aw / 2. / np.sqrt(2) / rho / gamma # 'natural' focusing wavenumber
  kbyn = aw / 2. / np.sqrt(2) / rho / gamma # 'natural' focusing wavenumber

else:
  kbxn = aw / 2. / np.sqrt(2) / rho / gamma # 'natural' focusing wavenumber
  kbyn = aw / 2. / np.sqrt(2) / rho / gamma # 'natural' focusing wavenumber
  awrms = aw / np.sqrt(2.) * np.sqrt(ux**2 + uy**2)

##################################################
# Sampling

if qFlatTopZ2[0] == 1:
    if qRoundZ2[0] == 1:
        lez2 = (7.5 * sigRoundZ2)  + sigz2     # flat top + gaussian
    else:
        lez2     = sigz2       # flat-top
        sigz2_in = 1.0e8
else:
    lez2     = 9.*sigz2     # gaussian
    sigz2_in = sigz2

lwz2 = 50.0                # Total size of sampled field in z2
lsys_z2 = lwz2 + lez2[0]      # Total length of sampled system in z2

dz2 = lambda_z2 / elms_per_wave    # Node spacing in z2
NNodesZ2 = lsys_z2 / dz2 + 1       # Number of radiation field nodes

dz2e = lambda_z2 / emps_per_wave   # Macroparticle spacing in z2
NMElecsZ2 = lez2 / dz2e            # Number of macroparticles in z2
NMElecsP2 = 19                     # Number of macroparticles to sample energy spread

dz = lambda_z2 / steps_per_per     # Step size in zbar
Nsteps = zbarprop / dz             # Number of steps

###################################################
# Chirp - 1% per sigma_z

dcgamma = 0.0 # 0.01 * gamma;   # Change in gamma per sigma_z due to chirp
chirp = -dcgamma / sigz;   # Energy chirp in z
chirpz2 = Lc * chirp;    # Energy chirp in z2

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

# Set up the beam and seed classes, which will contain
# the data to write in the beam and seed files

#beam = [bm() for ib in range(nbeams)] # list of electron beams

#seeds = [sd() for ic in range(nseeds)] # list of seeds

# Assign electron pulse data to beams

sigx = sigx / (sqrt(Lg) * sqrt(Lc))
sigy = sigy / (sqrt(Lg) * sqrt(Lc))

lex = 6.*sigx
ley = 6.*sigy

beam = bm()
beam.nbeams = nbeams
beam.bftype = 'simple'

for ij in range(nbeams):
  if qFlatTopZ2[ij] == 1:
      sigz2[ij] = 1e8          # Flat top case
  else:
      sigz2[ij] = sigz2[ij]
      lez2[ij] = sigz2[ij] * 8.

sigpx = np.array([1.])
sigpy = np.array([1.])
siggam = sig_gamma

lepx = 6. * sigpx
lepy = 6. * sigpy
legam = siggam * 7.

nmpx = np.array([10])
nmpy = np.array([10])
nmpz2 = NMElecsZ2

nmppx = np.array([10])
nmppy = np.array([10])
nmpgam = np.array([10])
qmatch = [False]
trloadmeth = (np.array([np.int_(2)]))
qequixy = [False]

bcenter = np.array([0.0])

#beam.sig = np.append(sigx, np.append(sigy, np.append(sigz2, np.append(sigpx, \
#                          np.append( sigpy, siggam) ) ) ))

beam.sig = np.concatenate( (sigx, sigy, sigz2, sigpx, sigpy, siggam) )
beam.le = np.concatenate( (lex, ley, lez2, lepx, lepy, legam) )
beam.nmps = np.int_(np.round(np.concatenate( (nmpx, nmpy, nmpz2, nmppx, nmppy, nmpgam) )))
beam.eratio = eratio
beam.emitx = emitx
beam.emity = emity
beam.alphax = alphax
beam.alphay = alphay
beam.chirp = chirp
beam.bcenz2 = bcenter
beam.Q = Q
beam.qRoundEj = qRoundZ2
beam.sigEj = sigRoundZ2
beam.bosc_Mag = bEnOscMag
beam.bosc_Fr  = bEnOscFr
beam.qmatch  = qmatch
beam.trload = trloadmeth

#beam.sig = sigx + sigy + sigz2 + sigpy + sigpy + siggam
#beam.le = lex + ley + lez2 + lepy + lepy + legam
#beam.nmps = nmpx + nmpy + nmpz2 + nmppx + nmppy + nmpgam
#beam.eratio = eratio
#beam.emitx = emitx
#beam.emity = emity
#beam.emitx = alphax
#beam.emity = alphay
#beam.chirp = chirp
#beam.bcenz2 = bcenter
#beam.Q = Q
#beam.qRoundEj = qRoundZ2
#beam.sigEj = sigRoundZ2
#beam.bosc_Mag = bEnOscMag
#beam.bosc_Fr  = bEnOscFr
#beam.qmatch  = qmatch
#beam.trload = trloadmeth

##################################################
# Field info

NNodesX = 1 # Next, create a class for the field vars
NNodesY = 1 # and a routine to write the seed file

lwx = 1E0
lwy = 1E0

lsys_x = lwx
lsys_y = lwy

filtFrac = 0.3
diffFrac = 1.0

#################################################
# Flags

qoned = True
qfieldevo = True
qEevolve  = True
qEFcouple = True
qFocusing = False
qMatchedBeam = False
qDiffraction = False
qFilter = True
qNoise = True
qDump = False
qResume = False
qStepFiles = True
qFormatFiles = True
qWriteZ = True
qWriteA = True
qWritePperp = True
qWritep2 = True
qWritez2 = True
qWritex = True
qWritey = True

########  OTHER DATA TEMP! #################

qScaled = True
qUndEnds = True
qFMesh = True

taper = 0.0
sZ0 = 0.0
kbetax_SF = 0.0
kbetay_SF = 0.0
wrFile = ''  # 'wr_cl.txt'
wrFreq_Full = 30
wrFreq_integrated = 30

################################################
# Write data

# Main input file:

f = open(inputfile, 'w')

# Header

f.write('! The main input parameters are described below - Puffin takes the namelist blocks at the\n')
f.write('! bottom of this file as input. This is the \'main\' input file, containing info about the\n')
f.write('! wiggler, field sampling, and general flags and other numerical instructions for the\n')
f.write('! simulation. This file also points to a beam file, describing the electron beam,\n')
f.write('! and optionally, a seed file (describing a radiation seed) and/or a lattice file.\n')
f.write('!\n')
f.write('!\n')
# Options

f.write('!                       OPTIONS\n')
f.write('!\n')
f.write('!\n')
f.write('! qOneD                      If TRUE, model 1D FEL, with only 1 node and 1 macroparticle in transverse dimensions\n')
f.write('! qFieldEvolve               if letting the radiation field evolve\n')
f.write('! qElectronsEvolve           if integrating electron equations\n')
f.write('! qElectronFieldCoupling     if allowing field to feedback onto the electron equations\n')
f.write('! qFocussing                 if strong intra-undulator focussing is included in the transverse plane\n')
f.write('! qMatchedBeam               if matching beam to undulator. If TRUE, electron pulse sigma and length in x,y,px,py are automatically calculated\n')
f.write('! qDiffraction               if modelling diffraction\n')
f.write('! qFilter                    TRUE to filter, if FALSE the low frequencies will just be ignored during diffraction\n')
f.write('! q_noise                    Shot noise in initial electron beam distribution\n')
f.write('! qResume                    If resuming from dump files left from a previous run\n')
f.write('! qMeasure                   If .TRUE., then use measured FFT plans in FFTW\n')
f.write('! qDumpEnd                   If .TRUE., the do full data dump at end of simulation\n')
f.write('! qInitWrLat\n')
f.write('!\n')
f.write('!\n')
f.write('!\n')
f.write('!\n')

# Field Sampling

f.write('!                  FIELD MESH DESCRIPTION\n')
f.write('!\n')
f.write('!\n')
f.write('! iNumNodesX               Number of nodes to sample radiation field in x direction\n')
f.write('! iNumNodesY               Number of nodes to sample radiation field in y direction\n')
f.write('! nodesPerLambdar            Number of nodes per resonant wavelength\n')
f.write('! sFModelLengthX             Length of radiation field model in x direction\n')
f.write('! sFModelLengthY             Length of radiation field model in y direction\n')
f.write('! sWigglerLengthZ2           Length of field model in z2-bar direction\n')
f.write('! iRedNodesX                 Length of central field section in x where electrons will not leave\n')
f.write('! iRedNodesY                 Length of central field section in y where electrons will not leave\n')
f.write('! sFiltFrac                  Specifies cutoff for high pass filter as fraction of resonant frequency - used in diffraction step\n')
f.write('! sDiffFrac                  Specifies diffraction step size as fraction of the undulator period\n')
f.write('! beta                       Absorption coefficient\n')

f.write('!\n')
f.write('!\n')

f.write('!                  INDEPENDANT VARIABLES\n')
f.write('!\n')
f.write('!\n')
f.write('! Scaled independent vars. These describe the reference beam and undulator \n')
f.write('! used for the scaling of the system. \n')
f.write('!\n')
f.write('!\n')
f.write('! srho       Pierce or FEL parameter, describing the strength of the interaction (or efficiency)\n')
f.write('! sux        Normalised magnitude of wiggler magnetic field x-vector: H=1 is helical, H=0 is planar\n')
f.write('! suy        Normalised magnitude of wiggler magnetic field y-vector: H=1 is helical, H=0 is planar\n')
f.write('! saw        peak undulator parameter\n')
f.write('! sgamma_r   Resonant, or reference, beam energy\n')
f.write('! lambda_w   Undulator period\n')

f.write('!')
f.write('!')
f.write('!                  SIMPLE UNDULATOR SETUP')
f.write('!')
f.write('! \'Quick\' or \'simple\' undulator. If not using a lattice file, then this \n')
f.write('! undulator will be used. If a lattice file is used then these parameters \n')
f.write('! are ignored.\n')
f.write('!')
f.write('!')
f.write('! zundType     Undulator type - \'curved\' , \'planepole\', \'helical\' or blank\n')
f.write('! taper        gradient of taper - d/dz of alpha\n')
f.write('! nPeriods           Number of wiggler periods\n')
f.write('! stepsPerPeriod     Number of integration steps per wiggler period\n')
f.write('!\n')
f.write('!\n')
f.write('!                  EXTERNAL FILES\n')
f.write('!\n')
f.write('! The beam file is required, if not specified then the beam file is\n')
f.write('! assumed = \'beam_file.in\', which must be in working directory.\n')
f.write('!\n')
f.write('! beam_file     Name of the beam file\n')
f.write('! seed_file     Name of the seed file\n')
f.write('! lattFile      Name of lattice file (optional).\n')
f.write('! wr_file       Name of the write file (optional).\n')

f.write('!\n')
f.write('!\n')
f.write('!               DATA WRITING FREQUENCY\n')
f.write('!\n')
f.write('! iWriteNthSteps     Steps to write data at\n')
f.write('! iWriteIntNthSteps  Steps to write integrated data at\n')
f.write('! sZ0                Starting zbar position\n')

f.write('!\n')
f.write('!\n')
f.write('!  Additional Beam Loading Options... (rarely used!!!) \n')
f.write('!\n')
f.write('! sPEOut        Percentage of macroparticles to write out\n')
f.write('! sEThreshold   Below the threshold level(%) * the average of real electrons are removed(ignored)\n')

f.write('\n')
f.write('\n')
f.write('\n')
f.write('\n')

f.write('&MDATA\n')
f.write(' qScaled                = ' + torf(qScaled) + '\n')
f.write(' qOneD                  = ' + torf(qoned) + '\n')
f.write(' qFieldEvolve           = ' + torf(qfieldevo) + '\n')
f.write(' qElectronsEvolve       = ' + torf(qEevolve) + '\n')
f.write(' qElectronFieldCoupling = ' + torf(qEFcouple) + '\n')
f.write(' qFocussing             = ' + torf(qFocusing) + '\n')
f.write(' qDiffraction           = ' + torf(qDiffraction) + '\n')
f.write(' qUndEnds               = ' + torf(qUndEnds) + '\n')
f.write(' qFMesh_G               = ' + torf(qFMesh) + '\n')
f.write(' beam_file              = ' +  '\'' + beamfile + '\'' + '\n')
if (seedfile != ''):
  f.write(' seed_file              = ' + '\'' + seedfile + '\'' + '\n')
#f.write(' lattFile               = ' + '\'' + lattFile + '\'' + '\n')
f.write(' wr_file                = ' + '\'' + wrFile + '\'' + '\n')

# Field Mesh:

f.write(' iNumNodesX             = ' + '{:d}'.format(int(math.floor(NNodesX))) + '\n')
f.write(' iNumNodesY             = ' + '{:d}'.format(int(math.floor(NNodesY))) + '\n')
f.write(' nodesPerLambdar        = ' + '{:d}'.format(int(math.floor(elms_per_wave)))  + '\n')
f.write(' sFModelLengthX         = ' + '{:.15E}'.format(lsys_x) + '\n')
f.write(' sFModelLengthY         = ' + '{:.15E}'.format(lsys_y) + '\n')
f.write(' sFModelLengthZ2        = ' + '{:.15E}'.format(lsys_z2) + '\n')
f.write(' iRedNodesX             = ' + '{:d}'.format(1) + '\n')
f.write(' iRedNodesY             = ' + '{:d}'.format(1) + '\n')
f.write(' sFiltFrac              = ' + '{:.15E}'.format(filtFrac) + '\n')
f.write(' sDiffFrac              = ' + '{:.15E}'.format(diffFrac) + '\n')
f.write(' sBeta                  = ' + '{:.15E}'.format(0.1) + '\n')

# Scaling parameters - sets up the frame

f.write(' srho                   = ' + '{:.15E}'.format(rho) + '\n')
f.write(' saw                    = ' + '{:.15E}'.format(aw) + '\n')
f.write(' sgamma_r               = ' + '{:.15E}'.format(gamma) + '\n')
f.write(' lambda_w               = ' + '{:.15E}'.format(lambda_w) + '\n')
f.write(' zundType               = ' + '\'' + undtype + '\'' + '\n')

# 'base' wiggler - with lambda_w above. Ignored with a lattice file

f.write(' lambda_w               = ' + '{:.15E}'.format(lambda_w) + '\n')
f.write(' taper                  = ' + '{:.15E}'.format(taper) + '\n')
f.write(' sKBetaXSF              = ' + '{:.15E}'.format(kbetax_SF) + '\n')
f.write(' sKBetaYSF              = ' + '{:.15E}'.format(kbetay_SF) + '\n')
f.write(' nPeriods               = ' + '{:d}'.format(int(N_w)) + '\n')

# integration/write options
f.write(' stepsPerPeriod         = ' + '{:d}'.format(steps_per_per) + '\n')
f.write(' sZ0                    = ' + '{:.15E}'.format(sZ0) + '\n')
f.write(' iWriteNthSteps         = ' + '{:d}'.format(wrFreq_Full) + '\n')
f.write(' iWriteIntNthSteps      = ' + '{:d}'.format(wrFreq_integrated) + '\n')
f.write(' /\n')

f.close()

print 'Main data file written to: ' + inputfile

#
# beam file

f = open(beamfile, 'w')

# Write beam files

#for ib in range(nbeams):
#    beam[ib].wrbeam(f,ib+1)
beam.wrbeam(f,1)
f.close()

print 'Beam file written to: ' + beamfile

#
# seed file

if (seedfile != ''):

  seed = sd()
  f = open(seedfile, 'w')

# Write seed data

  for ic in range(nseeds):
      seed.wrseed(f,ic+1)

  f.close()

  print 'Radiation seed file written to: ' + seedfile


#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Author: Paul J. Robinson <pjrobinson@ucla.edu>
#         Qiming Sun <osirpt.sun@gmail.com>
#

'''
Vasp CHGCAR file format

See also
https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
'''

import sys
import collections
import time
import numpy
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf.pbc import gto as pbcgto
from pyscf.tools import cubegen

if sys.version_info >= (3,):
    unicode = str

RESOLUTION = cubegen.RESOLUTION
BOX_MARGIN = cubegen.BOX_MARGIN

def density(cell, outfile, dm, nx=60, ny=60, nz=60, resolution=RESOLUTION):
    '''Calculates electron density and write out in CHGCAR format.

    Args:
        cell : Mole or Cell object
            Mole or pbc Cell. If Mole object is given, the program will guess
            a cubic lattice for the molecule.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.

    Returns:
        No return value. This function outputs a VASP chgcarlike file
        (with phase if desired)...it can be opened in VESTA or VMD or
        many other softwares

    Examples:

        >>> # generates the first MO from the list of mo_coefficents
        >>> from pyscf.pbc import gto, scf
        >>> from pyscf.tools import chgcar
        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
        >>> mf = scf.RHF(cell).run()
        >>> chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1())

    '''
    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)

    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    rho = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        if isinstance(cell, pbcgto.cell.Cell):
            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
        else:
            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
        rho[ip0:ip1] = lib.einsum('pi,ij,pj->p', ao, dm, ao)
    rho = rho.reshape(nx,ny,nz)

    cc.write(rho, outfile)

def orbital(cell, outfile, coeff, nx=60, ny=60, nz=60, resolution=RESOLUTION):
    '''Calculate orbital value on real space grid and write out in
    CHGCAR format.

    Args:
        cell : Mole or Cell object
            Mole or pbc Cell. If Mole object is given, the program will guess
            a cubic lattice for the molecule.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.

    Returns:
        No return value. This function outputs a VASP chgcarlike file
        (with phase if desired)...it can be opened in VESTA or VMD or
        many other softwares

    Examples:

        >>> # generates the first MO from the list of mo_coefficents
        >>> from pyscf.pbc import gto, scf
        >>> from pyscf.tools import chgcar
        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
        >>> mf = scf.RHF(cell).run()
        >>> chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0])

    '''
    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)

    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    orb_on_grid = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        if isinstance(cell, pbcgto.cell.Cell):
            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
        else:
            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
        orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
    orb_on_grid = orb_on_grid.reshape(nx,ny,nz)

    cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')

class CHGCAR(cubegen.Cube):
    '''  Read-write of the Vasp CHGCAR files  '''
    def __init__(self, cell, nx=60, ny=60, nz=60, resolution=RESOLUTION,
                 margin=BOX_MARGIN):
        if not isinstance(cell, pbcgto.cell.Cell):
            coord = cell.atom_coords()
            box = numpy.max(coord,axis=0) - numpy.min(coord,axis=0) + margin*2
            boxorig = numpy.min(coord,axis=0) - margin
            if resolution is not None:
                nx, ny, nz = numpy.ceil(box / resolution).astype(int)
            self.box = numpy.diag(box)
            lib.logger.warn(cell, 'Molecular system is found. FFT-grid is not '
                            'available for Molecule. Lattice (in Bohr)\n'
                            '%s\nand FFT grids %s are applied.',
                            self.box, (nx,ny,nz))
            self.mol = cell
            cell = cell.view(pbcgto.Cell)
            if (isinstance(cell.unit, (str, unicode)) and
                cell.unit.startswith(('B','b','au','AU'))):
                cell.a = self.box
            else:
                cell.a = self.box * lib.param.BOHR
            ptr = cell._atm[:,gto.PTR_COORD]
            cell._env[ptr+0] = coord[:,0] - boxorig[0]
            cell._env[ptr+1] = coord[:,1] - boxorig[1]
            cell._env[ptr+2] = coord[:,2] - boxorig[2]

        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.cell = cell
        self.box = cell.lattice_vectors()
        self.boxorig = numpy.zeros(3)
        self.vol = cell.vol

    def get_coords(self) :
        """  Result: set of coordinates to compute a field which is to be stored
        in the file.
        """
        xs = numpy.arange(self.nx) * (1./self.nx)
        ys = numpy.arange(self.ny) * (1./self.ny)
        zs = numpy.arange(self.nz) * (1./self.nz)
        xyz = lib.cartesian_prod((xs, ys, zs))
        coords = numpy.dot(xyz, self.box)
        return numpy.asarray(coords, order='C')

    def write(self, field, fname, comment=None):
        """  Result: .vasp file with the field in the file fname.  """
        assert(field.ndim == 3)
        assert(field.shape == (self.nx, self.ny, self.nz))
        if comment is None:
            comment = 'VASP file: Electron density in real space (e/Bohr^3)  '

        cell = self.cell

        # See CHGCAR format https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
        # the value of (total density * volume) was dumped
        field = field * self.vol

        boxA = self.box * lib.param.BOHR
        atomList= [cell.atom_pure_symbol(i) for i in range(cell.natm)]
        Axyz = zip(atomList, cell.atom_coords().tolist())
        Axyz = sorted(Axyz, key = lambda x: x[0])
        swappedCoords = [(vec[1]+self.boxorig) * lib.param.BOHR for vec in Axyz]
        vaspAtomicInfo = collections.Counter([xyz[0] for xyz in Axyz])
        vaspAtomicInfo = sorted(vaspAtomicInfo.items())
        with open(fname, 'w') as f:
            f.write(comment)
            f.write('PySCF Version: %s  Date: %s\n' % (pyscf.__version__, time.ctime()))
            f.write('1.0000000000\n')
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[0,0],boxA[0,1],boxA[0,2]))
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[1,0],boxA[1,1],boxA[1,2]))
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[2,0],boxA[2,1],boxA[2,2]))
            f.write(''.join(['%5.3s'%atomN[0] for atomN in vaspAtomicInfo]) + '\n')
            f.write(''.join(['%5d'%atomN[1] for atomN in vaspAtomicInfo]) + '\n')
            f.write('Cartesian \n')
            for ia in range(cell.natm):
                f.write(' %14.8f %14.8f %14.8f\n' % tuple(swappedCoords[ia]))
            f.write('\n')
            f.write('%6.5s %6.5s %6.5s \n' % (self.nx,self.ny,self.nz))
            fmt = ' %14.8e '
            for iz in range(self.nz):
                for iy in range(self.ny):
                    f.write('\n')
                    for ix in range(self.nx):
                        f.write(fmt % field[ix,iy,iz])

    def read(self, chgcar_file):
        raise NotImplementedError

if __name__ == '__main__':
    from pyscf.pbc import scf
    from pyscf.tools import chgcar
    cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
    mf = scf.RHF(cell).run()
    chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1()) #makes total density
    chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0]) # makes mo#1 (sigma)
    chgcar.orbital(cell, 'h2_mo2.CHGCAR', mf.mo_coeff[:,1]) # makes mo#2 (sigma*)


""" Function that adds common arguments for all input formats to the argument parser handle. """

import os

import numpy as np

from wmpl.Trajectory.Trajectory import Trajectory
from wmpl.Trajectory.GuralTrajectory import GuralTrajectory
from wmpl.Utils.TrajConversions import J2000_JD, jd2Date, equatorialCoordPrecession_vect, raDec2AltAz_vect

class MeteorObservation(object):
    """ Container for meteor observations.

        The points in arrays are RA and Dec in J2000 epoch, in radians.

        Arguments:
            jdt_ref: [float] Reference Julian date when the relative time is t = 0s.
            station_id: [str] Station ID.
            latitude: [float] Latitude +N in radians.
            longitude: [float] Longitude +E in radians.
            height: [float] Elevation above sea level (MSL) in meters.
            fps: [float] Frames per second.

        Keyword arguments:
            ff_name: [str] Name of the originating FF file.

    """
    def __init__(self, jdt_ref, station_id, latitude, longitude, height, fps, ff_name=None):

        self.jdt_ref = jdt_ref
        self.station_id = station_id
        self.latitude = latitude
        self.longitude = longitude
        self.height = height
        self.fps = fps

        self.ff_name = ff_name

        self.frames = []
        self.time_data = []
        self.x_data = []
        self.y_data = []
        self.azim_data = []
        self.elev_data = []
        self.ra_data = []
        self.dec_data = []
        self.mag_data = []
        self.abs_mag_data = []

    def addPoint(self, frame_n, x, y, azim, elev, ra, dec, mag):
        """ Adds the measurement point to the meteor.

        Arguments:
            frame_n: [flaot] Frame number from the reference time.
            x: [float] X image coordinate.
            y: [float] X image coordinate.
            azim: [float] Azimuth, J2000 in degrees.
            elev: [float] Elevation angle, J2000 in degrees.
            ra: [float] Right ascension, J2000 in degrees.
            dec: [float] Declination, J2000 in degrees.
            mag: [float] Visual magnitude.

        """

        self.frames.append(frame_n)

        # Calculate the time in seconds w.r.t. to the reference JD
        point_time = float(frame_n)/self.fps

        self.time_data.append(point_time)

        self.x_data.append(x)
        self.y_data.append(y)

        # Angular coordinates converted to radians
        self.azim_data.append(np.radians(azim))
        self.elev_data.append(np.radians(elev))
        self.ra_data.append(np.radians(ra))
        self.dec_data.append(np.radians(dec))
        self.mag_data.append(mag)

    def finish(self):
        """ When the initialization is done, convert data lists to numpy arrays. """

        self.frames = np.array(self.frames)
        self.time_data = np.array(self.time_data)
        self.x_data = np.array(self.x_data)
        self.y_data = np.array(self.y_data)
        self.azim_data = np.array(self.azim_data)
        self.elev_data = np.array(self.elev_data)
        self.ra_data = np.array(self.ra_data)
        self.dec_data = np.array(self.dec_data)
        self.mag_data = np.array(self.mag_data)

        # Sort by frame
        temp_arr = np.c_[self.frames, self.time_data, self.x_data, self.y_data, self.azim_data, \
        self.elev_data, self.ra_data, self.dec_data, self.mag_data]
        temp_arr = temp_arr[np.argsort(temp_arr[:, 0])]
        self.frames, self.time_data, self.x_data, self.y_data, self.azim_data, self.elev_data, self.ra_data, \
            self.dec_data, self.mag_data = temp_arr.T

    def __repr__(self):

        out_str = ''

        out_str += 'Station ID = ' + str(self.station_id) + '\n'
        out_str += 'JD ref = {:f}'.format(self.jdt_ref) + '\n'
        out_str += 'DT ref = {:s}'.format(jd2Date(self.jdt_ref, \
            dt_obj=True).strftime("%Y/%m/%d-%H%M%S.%f")) + '\n'
        out_str += 'Lat = {:f}, Lon = {:f}, Ht = {:f} m'.format(np.degrees(self.latitude),
            np.degrees(self.longitude), self.height) + '\n'
        out_str += 'FPS = {:f}'.format(self.fps) + '\n'

        out_str += 'Points:\n'
        out_str += 'Time, X, Y, azimuth, elevation, RA, Dec, Mag:\n'

        for point_time, x, y, azim, elev, ra, dec, mag in zip(self.time_data, self.x_data, self.y_data, \
            self.azim_data, self.elev_data, self.ra_data, self.dec_data, self.mag_data):

            if mag is None:
                mag = 0

            out_str += '{:.4f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:+.2f}, {:.2f}\n'.format(point_time,\
                x, y, np.degrees(azim), np.degrees(elev), np.degrees(ra), np.degrees(dec), mag)

        return out_str

def prepareObservations(meteor_list):
    """ Takes a list of MeteorObservation objects, normalizes all data points to the same reference Julian
        date, precesses the observations from J2000 to the epoch of date.

    Arguments:
        meteor_list: [list] List of MeteorObservation objects

    Return:
        (jdt_ref, meteor_list):
            - jdt_ref: [float] reference Julian date for which t = 0
            - meteor_list: [list] A list a MeteorObservations whose time is normalized to jdt_ref, and are
                precessed to the epoch of date

    """

    if meteor_list:

        # The reference meteor is the one with the first time of the first frame
        ref_ind = np.argmin([met.jdt_ref + met.time_data[0]/86400.0 for met in meteor_list])
        tsec_delta = meteor_list[ref_ind].time_data[0]
        jdt_delta = tsec_delta/86400.0

        ### Normalize all times to the beginning of the first meteor

        # Apply the normalization to the reference meteor
        meteor_list[ref_ind].jdt_ref += jdt_delta
        meteor_list[ref_ind].time_data -= tsec_delta

        meteor_list_tcorr = []

        for i, meteor in enumerate(meteor_list):

            # Only correct non-reference meteors
            if i != ref_ind:

                # Calculate the difference between the reference and the current meteor
                jdt_diff = meteor.jdt_ref - meteor_list[ref_ind].jdt_ref
                tsec_diff = jdt_diff*86400.0

                # Normalize all meteor times to the same reference time
                meteor.jdt_ref -= jdt_diff
                meteor.time_data += tsec_diff

            meteor_list_tcorr.append(meteor)

        ######

        # The reference JD for all meteors is thus the reference JD of the first meteor
        jdt_ref = meteor_list_tcorr[ref_ind].jdt_ref

        ### Precess observations from J2000 to the epoch of date
        meteor_list_epoch_of_date = []
        for meteor in meteor_list_tcorr:

            jdt_ref_vect = np.zeros_like(meteor.ra_data) + jdt_ref

            # Precess from J2000 to the epoch of date
            ra_prec, dec_prec = equatorialCoordPrecession_vect(J2000_JD.days, jdt_ref_vect, meteor.ra_data,
                meteor.dec_data)

            meteor.ra_data = ra_prec
            meteor.dec_data = dec_prec

            # Convert preccesed Ra, Dec to altitude and azimuth
            meteor.azim_data, meteor.elev_data = raDec2AltAz_vect(meteor.ra_data, meteor.dec_data, jdt_ref,
                meteor.latitude, meteor.longitude)

            meteor_list_epoch_of_date.append(meteor)

        ######

        return jdt_ref, meteor_list_epoch_of_date

    else:
        return None, None

def solveTrajectoryGeneric(jdt_ref, meteor_list, dir_path, solver='original', **kwargs):
    """ Feed the list of meteors in the trajectory solver and run it.

    Arguments:
        jdt_ref: [float] Reference Julian date for all objects in meteor_list.
        meteor_list: [list] A list of MeteorObservation objects.
        dir_path: [str] Path to the data directory.

    Keyword arguments:
        solver: [str] Solver choice:
            - "original" is the Monte Carlo solver
            - "gural" is the Gural solver (through C++ bindings)
        **kwargs: Keyword arguments for the trajectory solver.

    """

    # Create name of output directory
    output_dir = os.path.join(dir_path, jd2Date(jdt_ref, dt_obj=True).strftime("%Y%m%d-%H%M%S.%f"))

    # Init the trajectory solver
    if solver == 'original':
        traj = Trajectory(jdt_ref, output_dir=output_dir, meastype=1, **kwargs)

    elif solver.lower().startswith('gural'):
        velmodel = solver.lower().strip('gural')
        if len(velmodel) == 1:
            velmodel = int(velmodel)
        else:
            velmodel = 0

        traj = GuralTrajectory(len(meteor_list), jdt_ref, velmodel=velmodel, meastype=1, verbose=1,
            output_dir=output_dir)

    else:
        print('No such solver:', solver)
        return

    # Add meteor observations to the solver
    for meteor in meteor_list:

        if solver == 'original':

            comment = ''
            if hasattr(meteor, "ff_name"):
                comment = meteor.ff_name

            traj.infillTrajectory(meteor.ra_data, meteor.dec_data, meteor.time_data, meteor.latitude,
                meteor.longitude, meteor.height, station_id=meteor.station_id, \
                magnitudes=meteor.mag_data, comment=comment)

        elif solver.lower().startswith('gural'):

            # Extract velocity model is given
            try:
                velmodel = int(solver[-1])

            except:
                # Default to the exponential model
                velmodel = 3

            traj.infillTrajectory(meteor.ra_data, meteor.dec_data, meteor.time_data, meteor.latitude,
                meteor.longitude, meteor.height)

    # Solve the trajectory
    traj = traj.run()

    return traj

def addSolverOptions(arg_parser, skip_velpart=False):
    """ Adds common arguments for all input formats to the argument parser handle. """

    arg_parser.add_argument('-s', '--solver', metavar='SOLVER', help="""Trajectory solver to use. \n
        - 'original' - Monte Carlo solver
        - 'gural0' - Gural constant velocity
        - 'gural1' - Gural linear deceleration
        - 'gural2' - Gural quadratic deceleration
        - 'gural3' - Gural exponential deceleration
         """, type=str, nargs='?', default='original')

    arg_parser.add_argument('-t', '--maxtoffset', metavar='MAX_TOFFSET', nargs=1, \
        help='Maximum time offset between the stations.', type=float)

    arg_parser.add_argument('-v', '--vinitht', metavar='V_INIT_HT', nargs=1, \
        help='The initial veloicty will be estimated as the average velocity above this height (in km). If not given, the initial velocity will be estimated using the sliding fit which can be controlled with the --velpart option.', \
        type=float)

    if not skip_velpart:
        arg_parser.add_argument('-p', '--velpart', metavar='VELOCITY_PART', \
            help='Fixed part from the beginning of the meteor on which the initial velocity estimation using the sliding fit will start. Default is 0.25 (25 percent), but for noisier data this might be bumped up to 0.5.', \
            type=float, default=0.25)

    arg_parser.add_argument('-d', '--disablemc', \
        help='Do not use the Monte Carlo solver, but only run the geometric solution.', action="store_true")

    arg_parser.add_argument('-e', '--notimefit', \
        help="Do not estimate timing and velocity together. The times are assumed to be fixed. A list of time offsets per station can be provided, e.g. \"'CA001A':0.42,'CA0005':-0.3\" for speciflying offsets of 0.42 and -0.3 seconds for stations CA001A and CA0005, respectively. Make sure there are no spaces between the arguments, although spaces are fine in the station name.", \
        type=str, nargs="?", default=True)

    arg_parser.add_argument('-r', '--mcruns', metavar="MC_RUNS", \
        help='Number of Monte Carlo runs.', type=int, default=100)

    arg_parser.add_argument('-u', '--uncertgeom', \
        help='Compute purely geometric uncertainties.', action="store_true")

    arg_parser.add_argument('-m', '--mcstd', metavar='MC_STDDEVS', \
        help='Standard deviations of noise to add to measurements during the Monte Carlo procedure. 1.0 by default.', \
        type=float, default=1.0)

    arg_parser.add_argument('-g', '--disablegravity', \
        help='Disable gravity compensation.', action="store_true")

    arg_parser.add_argument('-l', '--plotallspatial', \
        help='Plot a collection of plots showing the residuals vs. time, lenght and height.', \
        action="store_true")

    arg_parser.add_argument('-j', '--jacchia', \
        help='Show the Jacchia exponential deceleration fit on plots with the dynamics.', \
        action="store_true")

    arg_parser.add_argument('-i', '--imgformat', metavar='IMG_FORMAT', nargs=1, \
        help="Plot image format. 'png' by default, can be 'pdf', 'eps',... ", type=str, default='png')

    arg_parser.add_argument('-x', '--hideplots', \
        help="Don't show generated plots on the screen, just save them to disk.", action="store_true")

    return arg_parser


from torch.utils.data import Dataset, DataLoader
import os,glob,random,cv2,json,torch,util,cv2,h5py,math
import numpy as np
import config as cfg
from SMPL_layer import SMPL_Layer

class Surreal(Dataset):
    # Note for the surreal data
    # The 3D keyps (in world frame), after aligning
    # by pelvis, match the output of smpl(theta,beta),
    # where (theta,beta) are the smpl params given in
    # in the world frame
    def __init__(self,phase):
        self.dataPhase=phase #"train", "val", or "test"
        self.pathPrefix="/data/cmu/extracted/data/{}".format(self.dataPhase)
        self.normalize=True
        self.cropSize=cfg.surrealCropSize
        self.pixFormat="NCHW"
        self.cameraSpace=cfg.surrealCameraSpace
        self.loadData()

    def pathFilter(self,rgbP,depthP):
        s1=rgbP.split("rgb")
        s2=depthP.split("depth")
        check1=s1[0]==s2[0]
        s11=s1[1].split(".")
        s22=s2[1].split(".")
        check2=s11[0]==s22[0]
        if check1 and check2:
            return True
        else:
            return False

    def swapRightLeftPose(self,pose):
        swapInds = np.array([
                0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11, 15, 16, 17, 12, 13, 14, 18,
                19, 20, 24, 25, 26, 21, 22, 23, 27, 28, 29, 33, 34, 35, 30, 31, 32,
                36, 37, 38, 42, 43, 44, 39, 40, 41, 45, 46, 47, 51, 52, 53, 48, 49,
                50, 57, 58, 59, 54, 55, 56, 63, 64, 65, 60, 61, 62, 69, 70, 71, 66,
                67, 68
            ], np.int32)

        signFlip = np.array([
                    1, 1, 1, 1, -1, -1, 1, -1, -1, 1, 1, 1, 1, -1, -1, 1, -1, -1, 1, 1, 1, 1, -1, -1, 1, -1, -1, 1, 1, 1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1
                ], dtype=pose.dtype)

        newPose=np.take(pose,swapInds)*signFlip

        return newPose

    def swapRightLeftJoints(self,joints):
        assert joints.shape[1] == 24
        swapInds = np.array([0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11,10,12,14,13, 15, 17,16, 19,18,21,20,23,22], np.int32)
        jointsSwap=np.take(joints,swapInds,axis=1)
        return jointsSwap

    def loadData(self):
        self.keyps2d=[]
        self.keyps3d=[]
        self.smplShapeParams=[]
        self.smplPoseParams=[]
        self.rgbImages=[]
        self.depthImages=[]
        self.zrot=[]
        self.camLoc=[]
        self.gender=[]

        fp=h5py.File(cfg.pathToSurreal+self.pathPrefix+"/annot.h5","r")
        keyps2d=np.array(fp["gt2d"])
        keyps3d=np.array(fp["gt3d"])
        smplShapeParams=np.array(fp["shape"])
        smplPoseParams=np.array(fp["pose"])
        rgbImages=np.array(fp["imageNameRGB"])
        depthImages=np.array(fp["imageNameDepth"])
        camLoc=np.array(fp["camLoc"])
        zRot=np.array(fp["zRot"])
        gender=np.array(fp["gender"])
        nData=min(len(rgbImages),len(depthImages))

        for idx in range(1000): #range(nData):
            #idx=3764   #2586 #1987
            currRgbPath=rgbImages[idx]
            currDepthPath=depthImages[idx]
            ret=self.pathFilter(currRgbPath,currDepthPath)
            if ret:
                fullRgbImgPath=cfg.pathToSurreal+self.pathPrefix+currRgbPath
                fullDepthImgPath=cfg.pathToSurreal+self.pathPrefix+currDepthPath
                if (not os.path.isfile(fullRgbImgPath)) or (not os.path.isfile(fullDepthImgPath)):
                    continue
                if np.any(keyps2d[:,:,idx]<0):
                    continue
                # check if all zero depth, if so remove
                tmp=np.load(fullDepthImgPath)
                loc=tmp!=float(1e10)
                if np.min(tmp)==np.max(tmp):
                    continue
                if np.min(tmp[loc])==np.max(tmp[loc]):
                    continue
                if gender[idx][0]==1:
                    continue
                self.keyps2d.append(keyps2d[:,:,idx])
                self.keyps3d.append(keyps3d[:,:,idx])
                self.smplShapeParams.append(smplShapeParams[idx])
                self.smplPoseParams.append(smplPoseParams[idx])
                self.rgbImages.append(fullRgbImgPath)
                self.depthImages.append(fullDepthImgPath)
                self.zrot.append(zRot[idx,0])
                self.camLoc.append(camLoc[idx,:])
                self.gender.append(gender[idx])

        assert(len(self.keyps2d)==len(self.keyps3d))
        assert(len(self.keyps2d)==len(self.smplShapeParams))
        assert(len(self.keyps2d)==len(self.smplPoseParams))
        assert(len(self.keyps2d)==len(self.rgbImages))
        assert(len(self.keyps2d)==len(self.depthImages))
        assert(len(self.keyps2d)==len(self.zrot))
        assert(len(self.keyps2d)==len(self.camLoc))
        assert(len(self.keyps2d)==len(self.gender))

    def normalizeDepth(self,depthImg):
        loc=depthImg!=float(1e10)
        normDepth=depthImg
        normDepth[loc]=(depthImg[loc]-np.min(depthImg[loc]))/(np.max(depthImg[loc])-np.min(depthImg[loc]))
        normDepth[~loc]=0

        return normDepth

    def __len__(self):
        return len(self.rgbImages)

    def __getitem__(self,index):
        rgbImg=cv2.imread(self.rgbImages[index])
        keyps2d=self.keyps2d[index].copy()

        """# plot points on image and save
        imgVis=rgbImg
        for r in range(len(keyps2d[0])):
            currX=int(keyps2d[0][r])
            currY=int(keyps2d[1][r])
            cv2.circle(imgVis,(currX,currY),2,(0,0,255),4)
        cv2.imwrite("/mnt/data/tmp/tmp/tmpKeyps1.png",imgVis)"""

        keyps2d[0]=[x*(1.0*self.cropSize/rgbImg.shape[1]) for x in keyps2d[0]]
        keyps2d[1]=[x*(1.0*self.cropSize/rgbImg.shape[0]) for x in keyps2d[1]]
        keyps2d/=self.cropSize
        keyps2d=np.asarray(keyps2d)
        resImg=cv2.resize(rgbImg,(self.cropSize,self.cropSize),interpolation=cv2.INTER_CUBIC)
        resImg=util.convert_image_by_pixformat_normalize(resImg,self.pixFormat,self.normalize)
        #resImg=rgbImg

        depthImg=self.normalizeDepth(np.load(self.depthImages[index]))
        #depthImg=cv2.resize(depthImg,(self.cropSize,self.cropSize),interpolation=cv2.INTER_NEAREST)
        depthImg=cv2.resize(depthImg,(self.cropSize,self.cropSize))
        depthImgCh=np.zeros((3,depthImg.shape[0],depthImg.shape[1]))
        depthImgCh[0,:,:]=depthImg
        depthImgCh[1,:,:]=depthImg
        depthImgCh[2,:,:]=depthImg
        keyps3d=self.keyps3d[index].copy()
        keyps3d=np.asarray(keyps3d)
        keyps3dw=keyps3d.copy()

        shape,pose=self.smplShapeParams[index].copy(),self.smplPoseParams[index].copy()
        currCamLoc=self.camLoc[index].copy()
        currExtrinsic,currR,currT=util.getSurrealExtrinsic(np.expand_dims(np.transpose(currCamLoc),axis=1))
        currZRot=self.zrot[index].copy()
        RzBody=np.array(((math.cos(currZRot),-math.sin(currZRot),0),(math.sin(currZRot),math.cos(currZRot),0),(0,0,1)))
        pose[0:3]=util.rotateBodyForVisSurreal(RzBody,pose[0:3])

        ############ swap l/r pose, flip y/z
        pose=self.swapRightLeftPose(pose)
        pose[1]=-pose[1]
        pose[2]=-pose[2]
        ############

        theta=np.concatenate((pose,shape),axis=0)

        ############ Swap keyps l/r
        keyps2d=self.swapRightLeftJoints(keyps2d)
        keyps3d=self.swapRightLeftJoints(keyps3d)
        keyps3dw=keyps3d.copy()
        ############

        if self.cameraSpace:
            # move 3d joints to camera space
            keyps3d=np.concatenate([keyps3d.transpose(),np.ones((keyps3d.shape[1],1))],axis=1).transpose()
            keyps3d=np.dot(currExtrinsic,keyps3d)
        pelvis=keyps3d[:,0].copy()
        keyps3d-=np.expand_dims(pelvis,1)

        currGender=self.gender[index].copy()
        currGender=currGender[0]

        return {
            "rgbImage":torch.from_numpy(resImg).float(),
            "depthImage":torch.from_numpy(depthImgCh).float(),
            "kp_2d":torch.from_numpy(keyps2d).float(),
            "kp_3d":torch.from_numpy(keyps3d).float(),
            "pelvis":torch.from_numpy(np.expand_dims(pelvis,1)).float(),
            "kp_3d_world":torch.from_numpy(np.asarray(keyps3dw)).float(),
            "theta":torch.from_numpy(theta).float(),
            "zrot":torch.from_numpy(np.asarray(currZRot)).float(),
            "camLoc":torch.from_numpy(np.asarray(currCamLoc)).float(),
            "extrinsic":torch.from_numpy(np.asarray(currExtrinsic)).float(),
            "intrinsic":torch.from_numpy(np.asarray(cfg.surrealIntrinsic)).float(),
            "rgbImageName":self.rgbImages[index],
            "depthImageName":self.depthImages[index],
            "gender":currGender,
            "data_set":"Surreal"
        }

class Human36M(Dataset):
    def __init__(self):
        self.onlySinglePerson=False
        self.normalize=True
        self.minPtsRequired=7
        self.scaleRange=[1.05,1.3]
        self.maxIntersectRatio=0.5
        self.cropSize=224
        self.pixFormat="NCHW"
        self.loadData()

    def loadData(self):
        self.keyps2d=[]
        self.keyps3d=[]
        self.smplShapeParams=[]
        self.smplPoseParams=[]
        self.images=[]
        self.boxes=[]

        fp=h5py.File(os.path.join(cfg.pathToHuman36m,"data","annot-copy.h5"),"r")
        keyps2d=np.array(fp["gt2d"])
        keyps3d=np.array(fp["gt3d"])
        smplShapeParams=np.array(fp["shape"])
        smplPoseParams=np.array(fp["pose"])
        images=np.array(fp["imagename"])

        assert(len(keyps2d)==len(keyps3d))
        assert(len(keyps2d)==len(smplShapeParams))
        assert(len(keyps2d)==len(smplPoseParams))
        assert(len(keyps2d)==len(images))

        def isValid(pts):
            r=[]
            for pt in pts:
                if pt[2]!=0:
                    r.append(pt)
            return r

        for idx in cfg.imgsLoad: #range(cfg.nImgsLoad): #range(len(keyps2d)):
            fullImgPath=os.path.join(cfg.pathToHuman36m,"data")+images[idx].decode()
            if not os.path.isfile(fullImgPath):
                continue
            keyp2d=keyps2d[idx].reshape((-1,3))
            if np.sum(keyp2d[:,2])<self.minPtsRequired:
                continue
            lt,rb,v=util.calc_aabb(isValid(keyp2d))
            self.keyps2d.append(np.array(keyp2d.copy(),dtype=np.float))
            self.boxes.append((lt,rb))
            self.keyps3d.append(keyps3d[idx].copy().reshape(-1,3))
            self.smplShapeParams.append(smplShapeParams[idx].copy())
            self.smplPoseParams.append(smplPoseParams[idx].copy())
            self.images.append(fullImgPath)

    def __len__(self):
        return len(self.images)

    def __getitem__(self,index):
        imgPath=self.images[index]
        keyps2d=self.keyps2d[index].copy()
        # remove head and neck points
        keyps2d=keyps2d[0:12,:]
        bbox=self.boxes[index]
        keyps3d=self.keyps3d[index].copy()
        keyps3d=keyps3d[0:12,:]

        """# plot points on image and save
        imgVis=cv2.imread(imgPath)
        for r in range(keyps2d.shape[0]):
            currX=int(keyps2d[r,0])
            currY=int(keyps2d[r,1])
            cv2.circle(imgVis,(currX,currY),2,(0,0,255),4)
            cv2.putText(imgVis,str(r),(currX,currY),cv2.FONT_HERSHEY_SIMPLEX,0.2,(255,255,255))
        cv2.imwrite("/home/BaseCode/trainedModels/5/tmpFull.png",imgVis)"""

        #scale=np.random.rand(4)*(self.scaleRange[1]-self.scaleRange[0])+self.scaleRange[0]
        scale=1
        bboxImg,bboxKps=util.cut_image(imgPath,keyps2d,scale,bbox[0],bbox[1])
        bboxKps=bboxKps[:,0:2]
        bboxKps[:,0]*=(1.0*self.cropSize/bboxImg.shape[0])
        bboxKps[:,1]*=(1.0*self.cropSize/bboxImg.shape[1])
        resImg=cv2.resize(bboxImg,(self.cropSize,self.cropSize),interpolation=cv2.INTER_CUBIC)
        #bboxKps[:,:2]=2.0*bboxKps[:, :2]*ratio-1.0
        bboxKps=(2.0*bboxKps*1.0/self.cropSize) - 1.0
        #print(bboxKps)

        shape,pose=self.smplShapeParams[index],self.smplPoseParams[index]
        #theta=np.concatenate((np.zeros(3),pose,shape),axis = 0)
        theta=np.concatenate((pose,shape),axis = 0)

        return {
            "image":torch.from_numpy(util.convert_image_by_pixformat_normalize(resImg,self.pixFormat,self.normalize)).float(),
            "kp_2d":torch.from_numpy(bboxKps).float(),
            "kp_3d":torch.from_numpy(keyps3d).float(),
            "theta":torch.from_numpy(theta).float(),
            "image_name":self.images[index],
            "w_smpl":1.0,
            "w_3d":1.0,
            "data_set":"Human3.6M"
        }

class MoSH(Dataset):
    def __init__(self):
        self.smpl=SMPL_Layer("neutral",cfg.smplModelPath)
        self.loadData()

    def loadData(self):
        fp=h5py.File(os.path.join(cfg.pathToMoSh,"data","mosh_gen","mosh_joints_annot.h5"),"r")
        self.smplShapeParams=np.array(fp["shape"])
        self.smplPoseParams=np.array(fp["pose"])
        self.smplJoints=np.array(fp["joints"])

        #self.smplShapeParams=self.smplShapeParams[:cfg.nSamplesXyz2Smpl,:]
        #self.smplPoseParams=self.smplPoseParams[:cfg.nSamplesXyz2Smpl,:]
        #self.smplJoints=self.smplJoints[:cfg.nSamplesXyz2Smpl,:]
        assert(len(self.smplShapeParams)==len(self.smplPoseParams))
        assert(len(self.smplShapeParams)==len(self.smplJoints))

    def __len__(self):
        return len(self.smplShapeParams)

    def __getitem__(self,index):
        shape,pose,joints=self.smplShapeParams[index],self.smplPoseParams[index],self.smplJoints[index]
        theta=np.concatenate((pose,shape),axis=0)
        pelv=joints[0,:]
        currXYZ=joints-pelv
        currXYZ=currXYZ[1:,:]
        return {
            "theta":torch.from_numpy(theta).float(),
            "xyz":torch.from_numpy(currXYZ).float(),
            "pelvis":torch.from_numpy(pelv).float(),
            "data_set":"MoSH"
        }

class COCO2017(Dataset):
    def __init__(self,datasetPath):
        self.onlySinglePerson=False
        self.normalize=True
        self.minPtsRequired=7
        self.scaleRange=[1.05,1.3]
        self.maxIntersectRatio=0.5
        self.cropSize=224
        self.pixFormat="NCHW"
        self.loadData(datasetPath)

    def convertToLsp14Pts(self,cocoPts):
        kpMap=[15,13,11,10,12,14,9,7,5,4,6,8,0,0]
        kpMap=[16,14,12,11,13,15,10,8,6,5,7,9,0,0]
        kps=np.array(cocoPts, dtype = np.float).reshape(-1, 3)[kpMap].copy()
        kps[12:,2]=0.0 #no neck, top head
        kps[:,2]/=2.0
        return kps

    def loadData(self,datasetPath):
        self.images=[]
        self.keyps=[]
        self.boxes=[]
        with open(os.path.join(datasetPath,"annotations","person_keypoints_train2017.json"),"r") as f:
            annotations=json.load(f)
        f.close()

        imgsIdInfo={}
        for inf in annotations["images"]:
            imgId=inf["id"]
            imgName=inf["file_name"]
            currAnno={}
            currAnno["image_path"]=os.path.join(datasetPath,"images","train-valid2017",imgName)
            currAnno["kps"]=[]
            currAnno["box"]=[]
            assert not (imgId in imgsIdInfo)
            imgsIdInfo[imgId]=currAnno

        for inf in annotations["annotations"]:
            imgId=inf["image_id"]
            kps=inf["keypoints"]
            boxInfo=inf["bbox"]
            box=[np.array([int(boxInfo[0]),int(boxInfo[1])]),np.array([int(boxInfo[0]+boxInfo[2]),int(boxInfo[1]+boxInfo[3])])]
            assert imgId in imgsIdInfo
            anno=imgsIdInfo[imgId]
            anno["box"].append(box)
            anno["kps"].append(self.convertToLsp14Pts(kps))

        self.createData(imgsIdInfo)

    def createData(self,imgsIdInfo):
        def checkCurrBboxOverlap(currBox,allOtherBoxes):
            for box in allOtherBoxes:
                if util.get_rectangle_intersect_ratio(currBox[0],currBox[1],box[0],box[1])>self.maxIntersectRatio:
                    return True
            return False
        for key,value in imgsIdInfo.items():
            currPath=value["image_path"]
            currKpsSet=value["kps"]
            currBoxSet=value["box"]
            if len(currBoxSet)>1:
                if self.onlySinglePerson:
                    continue

            for r in range(len(currBoxSet)):
                currKeyp=currKpsSet[r]
                currBox=currBoxSet[r]
                if np.sum(currKeyp[:,2])<self.minPtsRequired:
                    continue
                allOtherBoxes=currBoxSet.copy()
                allOtherBoxes.pop(r)
                if checkCurrBboxOverlap(currBox,allOtherBoxes):
                    continue
                self.images.append(currPath)
                self.keyps.append(currKeyp.copy())
                self.boxes.append(currBox.copy())

    def __len__(self):
        return len(self.images)

    def __getitem__(self,index):
        imgPath=self.images[index]
        keyps=self.keyps[index].copy()
        bbox=self.boxes[index]

        scale = np.random.rand(4)*(self.scaleRange[1]-self.scaleRange[0])+self.scaleRange[0]
        bboxImg,bboxKps=util.cut_image(imgPath,keyps,scale,bbox[0],bbox[1])
        ratio=1.0*self.cropSize/bboxImg.shape[0]
        bboxKps[:,:2]*=ratio
        resImg=cv2.resize(bboxImg,(self.cropSize,self.cropSize),interpolation=cv2.INTER_CUBIC)
        ratio=1.0/self.cropSize
        bboxKps[:,:2]=2.0*bboxKps[:, :2]*ratio-1.0

        return {
            "image":torch.tensor(util.convert_image_by_pixformat_normalize(resImg,self.pixFormat,self.normalize)).float(),
            "kp_2d":torch.tensor(bboxKps).float(),
            "image_name":imgPath,
            "data_set":"COCO 2017"
        }


#!/usr/bin/env python
# coding: utf8
#
# Copyright (c) 2021 Centre National d'Etudes Spatiales (CNES).
#
# This file is part of PANDORA_MCCNN
#
#     https://github.com/CNES/Pandora_MCCNN
#
# 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.
#
"""
This module contains functions to test mc-cnn fast and accurate
"""

import numpy as np
import torch
import torch.nn as nn

from mc_cnn.model.mc_cnn_fast import FastMcCnn
from mc_cnn.model.mc_cnn_accurate import AccMcCnnInfer

def point_interval(ref_features, sec_features, disp):
    """
    Computes the range of points over which the similarity measure will be applied

    :param ref_features: reference features
    :type ref_features: Tensor of shape (64, row, col)
    :param sec_features: secondary features
    :type sec_features: Tensor of shape (64, row, col)
    :param disp: current disparity
    :type disp: float
    :return: the range of the reference and secondary image over which the similarity measure will be applied
    :rtype: tuple
    """
    _, _, nx_ref = ref_features.shape
    _, _, nx_sec = sec_features.shape

    # range in the reference image
    left = (max(0 - disp, 0), min(nx_ref - disp, nx_ref))
    # range in the secondary image
    right = (max(0 + disp, 0), min(nx_sec + disp, nx_sec))

    return left, right

def run_mc_cnn_fast(img_ref, img_sec, disp_min, disp_max, model_path):
    """
    Computes the cost volume for a pair of images with mc-cnn fast

    :param img_ref: reference Dataset image
    :type img_ref:
    xarray.Dataset containing :
        - im : 2D (row, col) xarray.DataArray
    :param img_sec: secondary Dataset image
    :type img_sec:
    xarray.Dataset containing :
        - im : 2D (row, col) xarray.DataArray
    :param disp_min: minimum disparity
    :type disp_min: int
    :param disp_max: maximum disparity
    :type disp_max: int
    :param model_path: path to the trained network
    :type model_path: string
    :return: the cost volume ( similarity score is converted to a matching cost )
    :rtype: 3D np.array (row, col, disp)
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Create the network
    net = FastMcCnn()
    # Load the network
    net.load_state_dict(torch.load(model_path, map_location=device)["model"])
    net.to(device)
    net.eval()

    # Normalize images
    ref = img_ref["im"].copy(deep=True).data
    ref = (ref - ref.mean()) / ref.std()

    sec = img_sec["im"].copy(deep=True).data
    sec = (sec - sec.mean()) / sec.std()

    # Extracts the image features by propagating the images in the mc_cnn fast network
    # Right and left features of shape : (64, row-10, col-10)
    ref_features = net(torch.from_numpy(ref).to(device=device, dtype=torch.float), training=False)
    sec_features = net(torch.from_numpy(sec).to(device=device, dtype=torch.float), training=False)

    cv = computes_cost_volume_mc_cnn_fast(ref_features, sec_features, disp_min, disp_max)

    return cv

def computes_cost_volume_mc_cnn_fast(ref_features, sec_features, disp_min, disp_max):
    """
    Computes the cost volume using the reference and secondary features computing by mc_cnn fast

    :param ref_features: reference features
    :type ref_features: Tensor of shape (64, row, col)
    :param sec_features: secondary features
    :type sec_features: Tensor of shape (64, row, col)
    :return: the cost volume ( similarity score is converted to a matching cost )
    :rtype: 3D np.array (row, col, disp)
    """
    # Construct the cost volume
    disparity_range = np.arange(disp_min, disp_max + 1)

    # Allocate the numpy cost volume cv = (disp, col, row), for efficient memory management
    cv = np.zeros((len(disparity_range), ref_features.shape[2], ref_features.shape[1]), dtype=np.float32)
    cv += np.nan

    cos = nn.CosineSimilarity(dim=0, eps=1e-6)

    for disp in disparity_range:
        # Columns range in left and right image
        left, right = point_interval(ref_features, sec_features, disp)
        ind_d = int(disp - disp_min)
        cv[ind_d, left[0] : left[1], :] = np.swapaxes(
            (
                cos(ref_features[:, :, left[0] : left[1]], sec_features[:, :, right[0] : right[1]])
                .cpu()
                .detach()
                .numpy()
            ),
            0,
            1,
        )

    # Releases cache memory
    torch.cuda.empty_cache()

    # The minus sign converts the similarity score to a matching cost
    cv *= -1

    return np.swapaxes(cv, 0, 2)

def run_mc_cnn_accurate(img_ref, img_sec, disp_min, disp_max, model_path):
    """
    Computes the cost volume for a pair of images with mc-cnn accurate

    :param img_ref: reference Dataset image
    :type img_ref:
    xarray.Dataset containing :
        - im : 2D (row, col) xarray.DataArray
    :param img_sec: secondary Dataset image
    :type img_sec:
    xarray.Dataset containing :
        - im : 2D (row, col) xarray.DataArray
    :param disp_min: minimum disparity
    :type disp_min: int
    :param disp_max: maximum disparity
    :type disp_max: int
    :param model_path: path to the trained network
    :type model_path: string
    :return: the cost volume ( similarity score is converted to a matching cost )
    :rtype: 3D np.array (row, col, disp)
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Create the network
    net = AccMcCnnInfer()
    # Load the network
    net.load_state_dict(torch.load(model_path, map_location=device)["model"])
    net.to(device)
    net.eval()

    # Normalize images
    ref = img_ref["im"].copy(deep=True).data
    ref = (ref - ref.mean()) / ref.std()

    sec = img_sec["im"].copy(deep=True).data
    sec = (sec - sec.mean()) / sec.std()

    cv = net(
        torch.from_numpy(ref).to(device=device, dtype=torch.float),
        torch.from_numpy(sec).to(device=device, dtype=torch.float),
        disp_min,
        disp_max,
    )
    # Releases cache memory
    torch.cuda.empty_cache()
    return cv


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

# raspir
# Contact details: pust.marie-madlen@mh-hannover.de
# Last updated: 07 April 2021

# import
from __future__ import print_function
## batteries
import os
import sys
import argparse
import logging
import itertools
import random
from itertools import count, takewhile
## 3rd party
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy import fftpack, stats
from scipy.stats import linregress

# matplotlib init
matplotlib.use('Agg')

# logging
logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.DEBUG)
logging.getLogger('matplotlib.font_manager').disabled = True

# argparse#
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                      argparse.RawDescriptionHelpFormatter):
    pass

desc = 'raspir: rare species identifier'
epi = """DESCRIPTION:
Input: csv file (see raspir repo README)
Output:
* $PREFIX_[organism]_freq.png
  * graph: spectrum ~ frequence_per_cycle
* $PREFIX_final_stats.csv
  * csv table of final results (see raspir repo README)
"""
parser = argparse.ArgumentParser(description=desc, epilog=epi,
                                 formatter_class=CustomFormatter)
argparse.ArgumentDefaultsHelpFormatter
parser.add_argument('csv_file', metavar='csv_file', type=str,
                    help='Input csv file')
parser.add_argument('out_prefix', metavar='out_prefix', type=str,
                    help='output file prefix')
parser.add_argument('-e', '--error', type=float, default=0.01,
                    help='std-error cutoff parameter')
parser.add_argument('-a', '--alpha', type=float, default=0.05,
                    help='alpha parameter')
parser.add_argument('--version', action='version', version='0.0.1')

# Global parameters
norm_cpm = 1000000

# functions
def frange(start, stop, step):
    """
    Define range of the reference's read distance
    """
    return takewhile(lambda x: x < stop, count(start, step))

def read_count(x, min_reads=4):
    """
    Read count table
    """
    items_diff = [abs(j - i) for i, j in zip(x['Position'], x['Position'][1:])]
    a = [0]
    items_diff = a + items_diff
    x['items_diff'] = items_diff
    x['gap'] = np.where(x['items_diff'] == 1, 0, 1)
    x['readCount'] = np.cumsum(x['gap'])
    x = x.drop_duplicates(subset='readCount', keep='first')
    x = x.drop(['items_diff', 'gap'], 1)

    # remove all Organisms with less than `min_read` reads
    x['indexNames'] = x['readCount'].sum()
    store_index = x[x['indexNames'] < int(min_reads)].index
    x.drop(store_index, inplace=True)
    return x

def normalise_genome_position(x):
    """
    Normalise position (circular genome)
    """
    x['PositionNorm0'] = np.where(x['Position'] > (x['GenomeLength'] / 2),
                                  (x['GenomeLength'] - x['Position']),
                                  x['Position'])
    x['PositionNorm'] = x['PositionNorm0']**(1/2)

    # Reference position
    n_reads = x['readCount'].max()
    start_position_ref = int(1)
    end_position_ref = x['GenomeLength'].iloc[0]
    end_position_ref = end_position_ref + n_reads
    increase_by = (end_position_ref / n_reads)
    x['ref_Position'] = list(frange(start_position_ref, end_position_ref,
                                    increase_by))
    x['ref_Position'] = x['ref_Position'].astype(int)
    x['PositionNorm_ref0'] = np.where(x['ref_Position'] > (x['GenomeLength'] / 2),
                                      (x['GenomeLength'] - x['ref_Position']),
                                      x['ref_Position'])
    x['PositionNorm_ref'] = x['PositionNorm_ref0'].astype(int)
    return x

def make_time_domain(x):
    """
    Time domain signal
    """
    # check if read count of organism matches with minimum requirement
    mean_depth = int(x['Depth'].mean())
    n_reads = int(x['readCount'].max())
    species_name = x['Organism'].iloc[0]
    x['PositionNorm'] = x['PositionNorm'] * x['Depth']
    x['PositionNorm_ref'] = x['PositionNorm_ref'] * mean_depth

    reference_combinations_distances_sort = []
    real_combinations_distances_sort = []

    # calculate the biological distance
    real_read_positions = sorted(x['PositionNorm'])
    reference_read_positions = sorted(x['PositionNorm_ref'])

    if n_reads > int(1000):
        random.seed(222)
        real_select_random = random.sample(real_read_positions, 400)
        real_read_combinations_sub = list(itertools.combinations(real_select_random, 2))
        real_combinations_distances_sub = [abs(i - j) for i, j in real_read_combinations_sub]
        real_combinations_distances_sort_sub = sorted(real_combinations_distances_sub)
        real_combinations_distances_sort.append(real_combinations_distances_sort_sub)

        # calculate the reference distance
        reference_select_random = random.sample(reference_read_positions, 400)
        reference_read_combinations_sub = list(itertools.combinations(reference_select_random, 2))
        reference_combinations_distances_sub = [abs(i - j) for i, j in reference_read_combinations_sub]
        reference_combinations_distances_sort_sub = sorted(reference_combinations_distances_sub)
        reference_combinations_distances_sort.append(reference_combinations_distances_sort_sub)
    else:
        real_read_combinations = list(itertools.combinations(real_read_positions, 2))
        real_combinations_distances = [abs(i - j) for i, j in real_read_combinations]
        real_combinations_distances_sort1 = sorted(real_combinations_distances)
        real_combinations_distances_sort.append(real_combinations_distances_sort1)

        # calculate the reference distance
        reference_read_combinations = list(itertools.combinations(reference_read_positions, 2))
        reference_combinations_distances = [abs(i - j) for i, j in reference_read_combinations]
        reference_combinations_distances_sort1 = sorted(reference_combinations_distances)
        reference_combinations_distances_sort.append(reference_combinations_distances_sort1)

    # create output data frame
    df = pd.DataFrame(list(zip(reference_combinations_distances_sort, real_combinations_distances_sort)),
                      columns=['Reference', 'Real'])
    df['Organism'] = species_name
    df2 = df.apply(lambda i: i.explode() if i.name in ['Reference', 'Real'] else i)
    return df2

# frequency domain signal (fds)
def fourier_trans(x):
    species_name = x['Organism'].iloc[0]

    x['fft_ref1'] = np.fft.fft(x['Reference'])
    x['fft_bio1'] = np.fft.fft(x['Real'])

    x['fft_ref'] = [complex(np.around(items2.real), np.around(items2.imag)) for items2 in x['fft_ref1']]
    x['fft_bio'] = [complex(np.around(items2.real), np.around(items2.imag)) for items2 in x['fft_bio1']]

    x['fft_abs_ref'] = np.abs(x['fft_ref'])
    x['fft_abs_bio'] = np.abs(x['fft_bio'])
    x['fft_abs_ref_sqrt'] = np.around(x['fft_abs_ref'] / norm_cpm, 2)
    x['fft_abs_bio_sqrt'] = np.around(x['fft_abs_bio'] / norm_cpm, 2)

    # Pearson correlation
    if (sum(x['fft_abs_ref_sqrt']) > 0) & (sum(x['fft_abs_bio_sqrt']) > 0):
        pearson_corr = linregress(x['fft_abs_ref_sqrt'], x['fft_abs_bio_sqrt'])
        pearson_standard_error0 = pearson_corr[4]
        pearson_corr_r0, pearson_corr_p0 = stats.pearsonr(x['fft_abs_ref_sqrt'], x['fft_abs_bio_sqrt'])
        pearson_corr_r = round(pearson_corr_r0, 4)
        pearson_corr_p = round(pearson_corr_p0, 10)
        euclidean_dist_0 = np.linalg.norm(x['fft_abs_ref_sqrt']-x['fft_abs_bio_sqrt'])
        euclidean_dist = round(euclidean_dist_0, 1)
        pearson_standard_error = round(pearson_standard_error0, 5)
        return species_name, pearson_corr_r, pearson_corr_p, pearson_standard_error, euclidean_dist
    else:
        pearson_corr_r2 = 0
        pearson_corr_p2 = 0
        pearson_standard_error2 = 1
        euclidean_dist2 = 1
        return species_name, pearson_corr_r2, pearson_corr_p2, pearson_standard_error2, euclidean_dist2

def make_freq_images(x):
    logging.basicConfig(level=logging.ERROR)
    species_name = x['Organism'].iloc[0]
    path_real = x['PathName'].iloc[0]

    x['fft_ref1'] = np.fft.fft(x['Reference'])
    x['fft_bio1'] = np.fft.fft(x['Real'])

    x['fft_ref'] = [complex(np.around(items2.real), np.around(items2.imag)) for items2 in x['fft_ref1']]
    x['fft_bio'] = [complex(np.around(items2.real), np.around(items2.imag)) for items2 in x['fft_bio1']]

    x['fft_abs_ref'] = np.abs(x['fft_ref'])
    x['fft_abs_bio'] = np.abs(x['fft_bio'])
    x['fft_abs_ref_sqrt'] = np.around(x['fft_abs_ref'] / norm_cpm, 2)
    x['fft_abs_bio_sqrt'] = np.around(x['fft_abs_bio'] / norm_cpm, 2)

    # plot frequency signal
    val_bio = x['Real']
    val_ref = x['Reference']
    x_bio0 = x['fft_abs_bio_sqrt']
    x_reference0 = x['fft_abs_ref_sqrt']

    x_bio1 = x_bio0.sort_values()
    x_bio2 = pd.concat([x_bio1[::2], x_bio1[len(x_bio1)-2:0:-2]])
    x_bio = x_bio2.tolist()
    x_reference1 = x_reference0.sort_values()
    x_reference2 = pd.concat([x_reference1[::2], x_reference1[len(x_reference1)-2:0:-2]])
    x_reference = x_reference2.tolist()
    freqs_bio0 = fftpack.fftfreq(len(val_bio))
    freqs_bio = sorted(freqs_bio0)

    freqs_ref0 = fftpack.fftfreq(len(val_ref))
    freqs_ref = sorted(freqs_ref0)
    a = (len(freqs_ref) - len(x_reference))
    b = (len(freqs_bio) - len(x_bio))
    x_reference += [0]*a
    x_bio += [0]*b
    x_reference3 = np.sqrt(x_reference)
    sep = '_'

    fig, ax1 = plt.subplots(1, 1, figsize=(2.5, 2))
    ax1.plot(freqs_ref, x_reference3, "black", linewidth=1, linestyle='--', label="Reference", alpha=0.4)
    ax1.plot(freqs_bio, x_bio, 'blue', linewidth=1, linestyle='--', label="Sample")
    ax1.legend(framealpha=1, loc='upper right', fontsize=3)
    ax1.fill_between(freqs_ref, x_reference3, x_bio, facecolor='pink', alpha=0.2, interpolate=True)
    fig.text(0.5, 0.025, "Frequency per cycle",  ha='center', va='center', fontsize=4)
    fig.text(0.010, 0.5, "Spectrum", ha='center', va='center', rotation='vertical', fontsize=4)
    plt.xticks(fontsize=3)
    plt.yticks(fontsize=3)
    outfile = '_'.join([path_real, species_name, 'freq.png'])
    plt.savefig(outfile, dpi=600)
    logging.basicConfig(level=logging.DEBUG)
    logging.info('  File written: {}'.format(outfile))
    plt.close()

def final_table(x, set_error, set_alpha):
    """
    Create file data table
    """
    a0 = pd.DataFrame(x, columns=['Pearson'])
    a = a0.dropna(axis=0, how='all')
    a = a[['Species', 'r_value', 'p_value', 'stError', 'euclideanR0']] = pd.DataFrame(a['Pearson'].tolist())
    a.columns = ['a', 'b', 'c', 'd', 'e',
                 'Species', 'r_value', 'p_value', 'stError', 'euclideanR0']
    a.reset_index(drop=True, inplace=True)
    a = a.drop(a.columns[[0, 1]], axis=1)
    a = a[['Species', 'r_value', 'p_value', 'stError', 'euclideanR0']]
    a['euclidean'] = np.around((1 / a['euclideanR0']) * 1000, 3)
    a['distribution'] = np.where(
        (a['p_value'] < set_alpha) & (a['r_value'] > 0.5) & (a['stError'] < set_error) & (a['euclidean'] < 0.5),
        'uniform', 'nonuniform')
    b1 = a[['Species', 'r_value', 'p_value', 'stError', 'euclidean', 'distribution']]
    b2 = b1[b1.distribution == 'uniform']
    return b2

def process_csv(file_name, out_prefix, args):
    with open(file_name, newline='') as inF:
        df = pd.read_csv(inF, delimiter=',')

        # filtering human reads
        pattern_del = '1_1_1_'
        filter_approach = df['Organism'].str.contains(pattern_del, na=False)
        df = df[~filter_approach]
        logging.info('1) Human reads have been removed')

        # counting reads per organism
        df = df.dropna(subset=['GenomeLength'])
        df_filter1 = df.groupby('Organism').apply(read_count)
        logging.info('2) Continue with the first position of each read')

        if df_filter1.empty is True:
            logging.warning('Note: Dataset is empty')
            outfile = '{}_filtered.csv'.format(out_prefix)
            df_filter1.to_csv(outfile, index=False)
            logging.info('  File written: {}'.format(outfile))
            return None
        else:
            df_filter = df_filter1.reset_index(drop=True)
            df_filter2 = df_filter.groupby('Organism').apply(normalise_genome_position)
            logging.info('3) Genome position normalised')

            position_domain0 = df_filter2.reset_index(drop=True)
            position_domain = position_domain0.groupby("Organism").apply(make_time_domain)
            logging.info('4) Time-domain signal built')

            frequency_domain0 = position_domain.reset_index(drop=True)
            frequency_domain0['PathName'] = os.path.join(out_prefix)
            frequency_domain = frequency_domain0.groupby('Organism').apply(fourier_trans)
            logging.info('5) Frequency-domain signal generated')

            frequency_domain0.groupby('Organism').apply(make_freq_images)
            logging.info('6) Frequency plots produced')

            stat_table = final_table(frequency_domain,
                                     set_error=args.error,
                                     set_alpha=args.alpha)
            logging.info('7) Output table has been generated')

            outfile = '{}_final_stats.csv'.format(out_prefix)
            stat_table.to_csv(outfile, index=False)
            logging.info('  File written: {}'.format(outfile))
            logging.info('8) Run successful')

def main(args):
    """
    Main interface
    """
    # output directory
    outdir = os.path.split(args.out_prefix)[0]
    if outdir != '' and not os.path.isdir(outdir):
        os.makedirs(args.outdir)
    # processing each file
    process_csv(args.csv_file, out_prefix=args.out_prefix, args=args)

if __name__ == "__main__":
    args = parser.parse_args()
    main(args)


# Copyright 2020 Makani Technologies LLC
#
# 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.

"""Functions for working with state-space models."""

import copy

import numpy as np

# pylint doesn't like capital letters in variable names, in contrast
# to control systems conventions.
# pylint: disable=invalid-name

class SignalListInvalidArgumentException(Exception):
  """Raised if an invalid argument is passed to a method of SignaList."""
  pass

class SignalList(object):
  """Representation of a named list of signals.

  A signal list consists of a non-repeating ordered set of signal
  names (each representing a scalar valued signal).
  """

  def __init__(self, names):
    """Constructor.

    Args:
      names: List of signal names.

    Raises:
      SignalListInvalidArgumentException: If there are repeated names.
    """
    self.names = copy.copy(names)
    if not all([isinstance(name, str) for name in self.names]):
      raise SignalListInvalidArgumentException('Names must be strings.', names)
    # Test that there are not repeated names.
    if len(self.names) != len(set(self.names)):
      raise SignalListInvalidArgumentException('Repeated signal name.', names)
    self._names_to_indices = {
        name: i for i, name in enumerate(self.names)
    }

  def __repr__(self):
    return ', '.join(self.names)

  def __getitem__(self, indices):
    """Select a subset of an input list.

    Args:
      indices: Can be a slice or a list.  If a list, may contain
          a mixture of integer entries and string entries.

    Raises:
      SignalListInvalidArgumentException: If indices is not a list.

    Returns:
      A new signal list consisting of a subset of the existing list
      in the given order (units are tracked appropriately).
    """
    if isinstance(indices, slice):
      selection = range(len(self))[indices]
    elif not isinstance(indices, list):
      raise SignalListInvalidArgumentException(indices)
    else:
      selection = copy.copy(indices)

    for i in range(len(selection)):
      if isinstance(selection[i], int):
        selection[i] = self.names[selection[i]]

    return SignalList(selection)

  def __len__(self):
    return len(self.names)

  def __add__(self, other):
    return SignalList(self.names + other.names)

  def GetIndices(self, names):
    """Get the integer indices corresponding to a list of signal names."""
    return [self._names_to_indices[name] for name in names]

  def AddSuffix(self, suffix):
    return SignalList([name + suffix for name in self.names])

class SystemInvalidArgumentException(Exception):
  """Raised if an invalid argument is passed to a method of System."""
  pass

class SystemBadDimensionException(Exception):
  """Raised if there is a mismatch in dimensions of arguments."""
  pass

class SystemSamplePeriodMismatchException(Exception):
  """Raised if there is a mismatch in sample periods."""
  pass

class SystemBadSamplePeriodException(Exception):
  """Raised if an illegal sample period is given."""
  pass

class SystemIllPosedException(Exception):
  """Raised if an ill-posed feedback loop is closed."""
  pass

class System(object):
  """Class for representing as state-space model with named inputs."""

  def __init__(self, A, B, C, D, Ts, state_list, input_list, output_list):
    """Constructor for a new system.

    If A, B, or C are empty, then all three must be empty and
    the system represents a constant gain matrix determined
    by D.

    If A is non-empty and D is empty, then D is taken to be all zeros.

    Args:
      A: nx-by-nx matrix.
      B: nx-by-nu matrix.
      C: ny-by-nx matrix.
      D: ny-by-nu matrix.
      Ts: Sample period (zero indicates continuous time models, -1.0
          indicates a DT model with unspecified sample period).
      state_list: SignalList of length nx.
      input_list: SignalList of length nu.
      output_list: SignalList of length ny.

    Raises:
      SystemBadSamplePeriodException: If Ts has an invalid value.
      SystemBadDimensionException: If the arguments have inconsistent
          dimensions.
    """
    if not state_list:
      state_list = SignalList([])
    if isinstance(state_list, list):
      state_list = SignalList(state_list)
    if isinstance(input_list, list):
      input_list = SignalList(input_list)
    if isinstance(output_list, list):
      output_list = SignalList(output_list)
    if Ts < 0.0 and Ts != -1.0:
      raise SystemBadSamplePeriodException()
    self.Ts = Ts

    self._A = np.matrix(A)
    self._B = np.matrix(B)
    self._C = np.matrix(C)
    self._D = np.matrix(D)
    if self._A.size == 0:
      if self._B.size > 0 or self._C.size > 0:
        raise SystemBadDimensionException(A, B, C, D)
      self.nx = 0
      self.ny = self._D.shape[0]
      self.nu = self._D.shape[1]
      self._A = np.matrix(np.zeros((self.nx, self.nx)))
      self._B = np.matrix(np.zeros((self.nx, self.nu)))
      self._C = np.matrix(np.zeros((self.ny, self.nx)))
    else:
      self.nx = self._A.shape[0]
      self.nu = self._B.shape[1]
      self.ny = self._C.shape[0]
    if self._D.size == 0:
      self._D = np.matrix(np.zeros((self.ny, self.nu)))

    if (not np.array_equal(self._A.shape, [self.nx, self.nx])
        or not np.array_equal(self._B.shape, [self.nx, self.nu])
        or not np.array_equal(self._C.shape, [self.ny, self.nx])
        or not np.array_equal(self._D.shape, [self.ny, self.nu])):
      raise SystemBadDimensionException(self._A.shape, self._B.shape,
                                        self._C.shape, self._D.shape)

    if not isinstance(state_list, SignalList) or len(state_list) != self.nx:
      raise SystemBadDimensionException('Bad state list.', state_list)

    if not isinstance(input_list, SignalList) or len(input_list) != self.nu:
      raise SystemBadDimensionException('Bad input list.', input_list)

    if not isinstance(output_list, SignalList) or len(output_list) != self.ny:
      raise SystemBadDimensionException('Bad output list.', output_list)

    self.states = state_list
    self.inputs = input_list
    self.outputs = output_list

  def __repr__(self):
    return ('%d-by-%d state-space model (Ts = %g) with %d states.' %
            (self.ny, self.nu, self.Ts, self.nx)
            + '\nInputs: ' + self.inputs.__repr__()
            + '\nOutputs: ' + self.outputs.__repr__()
            + '\nA:\n' + str(self._A) + '\nB:\n' + str(self._B)
            + '\nC:\n' + str(self._C) + '\nD:\n' + str(self._D))

  def __getitem__(self, io):
    """Construct a new system with a subset of the inputs and outputs.

    Args:
      io: A tuple of two lists of strings.  The first list selects a subset
        of the outputs of the system, the second selects a subset of the inputs.
        Either can be the trivial slice ':'.

    Raises:
      SystemInvalidArgumentException: If an invalid argument is passed.

    Returns:
      A new System object with the reduced input and output dimensions.
    """
    if not isinstance(io, tuple) or len(io) != 2:
      raise SystemInvalidArgumentException('Index must be a tuple.')

    if isinstance(io[1], slice):
      if io[1] != slice(None, None, None):
        raise SystemInvalidArgumentException('Only ":" slices are allowed.')
      input_names = [self.inputs.names[i] for i in range(self.nu)]
    else:
      input_names = io[1]

    if isinstance(io[0], slice):
      if io[0] != slice(None, None, None):
        raise SystemInvalidArgumentException('Only ":" slices are allowed.')
      output_names = [self.outputs.names[i] for i in range(self.ny)]
    else:
      output_names = io[0]

    (inputs, _, B, _) = self._PartitionInputs(self._B, input_names)
    (outputs, _, C, _) = self._PartitionOutputs(self._C, output_names)
    (_, _, D, _) = self._PartitionInputs(self._D, input_names)
    (_, _, D, _) = self._PartitionOutputs(D, output_names)

    return System(self._A, B, C, D, self.Ts, self.states, inputs, outputs)

  def _PartitionOutputs(self, matrix, output_names):
    """Partition the rows of an ny-by-m matrix.

    Args:
      matrix: An ny-by-m matrix.
      output_names: Names of the outputs to be included in the first
          part of the partition.

    Raises:
      SystemBadDimensionException: If matrix has the wrong number of rows.

    Returns:
      A tuple (outputs, other_outputs, output_matrix,
      other_output_matrix).  The first two entries are SignalLists
      containing the outputs in output_names and the other outputs.
      The first matrix is the sub-matrix corresponding to the
      output_names and the remaining rows are in other_output_matrix.
    """
    if matrix.shape[0] != self.ny:
      raise SystemBadDimensionException(matrix)

    outputs = self.outputs[output_names]
    output_indices = self.outputs.GetIndices(output_names)
    other_output_indices = [
        i for i in range(self.ny) if i not in output_indices
    ]
    other_outputs = self.outputs[other_output_indices]

    output_matrix = matrix[output_indices, :]
    other_output_matrix = matrix[other_output_indices, :]

    return (outputs, other_outputs, output_matrix, other_output_matrix)

  def _PartitionInputs(self, matrix, input_names):
    """Partition the columns of an m-by-nu matrix.

    Args:
      matrix: An m-by-nu matrix.
      input_names: Names of the inputs to be included in the first
          part of the partition.

    Raises:
      SystemBadDimensionException:  If matrix has the wrong number of columns.

    Returns:
      A tuple (inputs, other_inputs, input_matrix,
      other_input_matrix).  The first two entries are SignalLists
      containing the inputs in input_names and the other inputs.
      The first matrix is the sub-matrix corresponding to the
      input_names and the remaining rows are in other_input_matrix.
    """
    if matrix.shape[1] != self.nu:
      raise SystemBadDimensionException(matrix)

    inputs = self.inputs[input_names]
    input_indices = self.inputs.GetIndices(input_names)
    other_input_indices = [
        i for i in range(self.nu) if i not in input_indices
    ]
    other_inputs = self.inputs[other_input_indices]

    input_matrix = matrix[:, input_indices]
    other_input_matrix = matrix[:, other_input_indices]

    return (inputs, other_inputs, input_matrix, other_input_matrix)

  def GetStateSpaceModel(self):
    """Return the state-space description of the system."""
    return self._A, self._B, self._C, self._D, self.Ts

  def ReduceStates(self, state_names):
    """Returns a new system model that keeps only a subset of the states.

    Args:
      state_names: List of state names to retain.

    Returns:
      A new state-space model with only the desired states retained.
    """
    state_indices = self.states.GetIndices(state_names)
    states = self.states[state_names]

    A = self._A[[[s] for s in state_indices], state_indices]
    B = self._B[[s for s in state_indices], :]
    C = self._C[:, state_indices]

    return System(A, B, C, self._D, self.Ts, states, self.inputs, self.outputs)