import os
import numpy as np
from glob import iglob
import shutil
import os
import readline
import pprint
import sys
from sys import argv
#from NN_phasecorrection import check_phases

'''This file should interpolate between a starting geometry taken from a QM.in file and an end-geometry (at which NNs broke down), also taken from a QM.in file_end
      - The files should be transfered into zmat-files and linear interpolation between the start- and end-geometry will be carried out
      - Later, the geometries will be written into the form of a QM.in file, the interface will then generate QM.out files - those will be written into an output.dat format and the output_all.dat can be appended later
      - In the end, the phases should be compared between geometries and corrected in the output.dat files. After the last geometry is corrected, the
        calculation using a QM-interface, should be carried out with corrected phases
'''

#============================================================================================================================================================#
#======================================================================QM.out 2 output.dat===================================================================#
#============================================================================================================================================================#
def qm2outputdat(Properties,outputdatfile):
  #in order to use the generated data, the important quantities for running NNs (Energy, SOC, Dipole, NAC) will be written into an output.dat format
  #therefore, the header of the output_all.dat file will be taken  -  every matrix (e.g. U matrix, property matrix) will be filled with zeros
  #iterate over QM.out files
  #TODO maybe change Overlap to True again - but right now although overlap and write_overlap is given, no overlaps are written by molpro
  OUTPUT['Overlap']=False
  file=open("output_new.dat", 'w')
  nmstates = OUTPUT['nmstates']
  natoms = OUTPUT['natoms']
  string_output=header
  string_output+='! 0 Step\n 0\n'
  string_hamilton='! 1 Hamiltonian (MCH) in a.u.\n'
  string_umatrix='! 2 U matrix\n'
  string_dipole_x='! 3 Dipole moments X (MCH) in a.u.\n'
  string_dipole_y='! 3 Dipole moments Y (MCH) in a.u.\n'
  string_dipole_z='! 3 Dipole moments Z (MCH) in a.u.\n'
  string_overlap=''
  string_coefficient='! 5 Coefficients (diag)\n'
  string_hopping='! 6 Hopping Probabilities\n'
  string_ekin='! 7 Ekin (a.u.)\n0.000000000\n'
  string_states='! 8 states (diag, MCH)\n0.000000000 0.000000000\n'
  string_random='! 9 Random number\n0.000000000\n'
  string_runtime='! 10 Runtime (sec)\n0.000000000\n'
  #this is done because the file has problems to do the last step - thus one step more is done and an empty file will be written and removed later
  #if index<= nsteps-1:
  string_geom='! 11 Geometry in a.u.\n'
  string_velocities='! 12 Velocities in a.u.\n'
  string_property2d=''
  string_property1d=''
  string_gradient=''
  string_nac=''

  #write string for geometires
  for atomnr in range(natoms):
    string_geom+='%20.12f %20.12f %20.12f\n' %(Properties['Geometries'][atomnr][0],Properties['Geometries'][atomnr][1],Properties['Geometries'][atomnr][2])

  if OUTPUT['Overlap']==True:
    string_overlap+='! 4 Overlap matrix (MCH)\n'
    for numberofstates in range(nmstates):
      for real_imag in range(nmstates*2):
        string_overlap+='0.000000000 '
      string_overlap+='\n'

  for numberofstates in range(nmstates):
    string_coefficient+='0.0000000000 0.0000000000\n'
    string_hopping+='0.000000000 0.000000000\n'
    for real_imag in range(nmstates*2):
      #if phase correction applied: Properties should be replaced with phasecorrected properties. Hamiltonian --> Hamiltonian_phasecorrected
      if (2*numberofstates)==real_imag:
        string_hamilton+='%20.12f ' %(Properties['Hamiltonian'][numberofstates][real_imag]-OUTPUT['ezero'])
      else:
        string_hamilton+='%20.12f ' %(Properties['Hamiltonian'][numberofstates][real_imag])
      string_umatrix+='0.000000000 '
      string_dipole_x+='%20.12f ' %Properties['Dipole_x'][numberofstates][real_imag]
      string_dipole_y+='%20.12f ' %Properties['Dipole_y'][numberofstates][real_imag]
      string_dipole_z+='%20.12f ' %Properties['Dipole_z'][numberofstates][real_imag]
    string_umatrix+='\n'
    string_hamilton+='\n'
    string_dipole_x+='\n'
    string_dipole_y+='\n'
    string_dipole_z+='\n'
  for numberofatoms in range(natoms):
    string_velocities+='0.000000000 0.000000000 0.000000000\n'

  if OUTPUT['Gradient']==True:
    for numberofstates in range(nmstates):
      string_gradient+='! 15 Gradients (MCH) State %i\n' %(int(numberofstates)+1)
      xyz=0
      for numberofatoms in range(natoms*3):
        xyz+=1
        string_gradient+='%20.12f ' %Properties['Gradient'][numberofstates][numberofatoms]
        if xyz>=3:
          xyz=0
          string_gradient+='\n'

  if int(OUTPUT['Property2d'])==1:
    n_property2d=int(OUTPUT['n_property2d'])
    for number in range(n_property2d):
      string_property2d+='! 13 Property matrix (MCH)  %i : N/A\n' %(int(number)+1)
      for numberofstates in range(nmstates):
        for real_imag in range(nmstates*2):
          string_property2d+='0.000000000 '
        string_property2d+='\n'

  if int(OUTPUT['Property1d'])==1:
    n_property1d=int(OUTPUT['n_property1d'])
    for number in range(n_property1d):
      string_property1d+='! 14 Property vector (MCH)  %i : N/A\n' %(int(number)+1)
      for numberofstates in range(nmstates):
        string_property1d+='0.000000000\n'

  if OUTPUT['NACdr']==True:
    state=0
    firststate=1
    for nacindex in range(nmstates*nmstates):
      state+=1
      string_nac+='! 16 NACdr matrix element (MCH)  %i %i\n' %(firststate, state)
      if state >= nmstates:
        state=0
        firststate+=1
      xyz=0
      for numberofatoms_xyz in range(natoms*3):
        xyz+=1
        string_nac+='%20.12f ' %(Properties['NonAdiabaticCouplings'][nacindex][numberofatoms_xyz])
        if xyz>=3:
          xyz=0
          string_nac+='\n'

  string_output+=string_hamilton+string_umatrix+string_dipole_x+string_dipole_y+string_dipole_z+string_overlap+string_coefficient+string_hopping+string_ekin+string_states+string_random+string_runtime+string_geom+string_velocities+string_property2d+string_property1d+string_gradient+string_nac
  file.write(string_output)
  file.close()

  return OUTPUT

"""def generate_outputdat(scanpath, nsteps):
  #merge all output.dat files to one file called output.dat
  outputfile = open(scanpath+"/output.dat", "w")
  header_exists=False
  nsteps=int(nsteps)
  for index in range(1,nsteps):
    infile=open(scanpath+'/output'+str(index)+'.dat', 'r').readlines()
    is_header=True
    for line in infile:
      if not header_exists:
        outputfile.write(line)
      if 'End of header' in line:
        is_header=False
        header_exists=True
        continue
      if not is_header:
        outputfile.write(line)
  outputfile.close()
  for i in range(1,nsteps):
    os.system('rm %s/output%i.dat' %(scanpath,i))"""

#=================================================================GET HEADER OF OUTPUT.DAT===================================================================#
def get_header(outputdatfile):

  OUTPUT={ 'Overlap':      False,
           'Gradient':     False,
           'NACdr':        False,
           'Property1d':   0,
           'Property2d':   0,
           'n_property1d': 1,
           'n_property2d': 1}

  data=open(outputdatfile,'r').readlines()
  iline=-1
  header=''
  for line in data:
    iline+=1
    line_string=data[iline]
    line=line.strip()
    header+='%s' %line_string
    if line.startswith('nstates_m'):
      line = line.split()
      n_singlets=int(line[1])
      n_triplets=int(line[3])
      n_dubletts=int(line[2])
      nmstates=1*int(line[1])+2*int(line[2])+3*int(line[3])
      OUTPUT['nmstates']=int(nmstates)
      OUTPUT['n_singlets']=int(n_singlets)
      OUTPUT['n_triplets']=int(n_triplets)
    elif line.startswith('natom'):
      line = line.split()
      natoms=line[1]
      OUTPUT['natoms']=int(natoms)
    elif line.startswith('write_overlap'):
      line=line.split()
      if int(line[1])==1:
        OUTPUT['Overlap']=True
    elif line.startswith('write_grad'):
      line=line.split()
      if int(line[1])==1:
        OUTPUT['Gradient']=True
    elif line.startswith('write_nacdr'):
      line=line.split()
      if int(line[1])==1:
        OUTPUT['NACdr']=True
    elif line.startswith('write_property1d'):
      line=line.split()
      if int(line[1])==1:
        OUTPUT['Property1d']=1
    elif line.startswith('write_property2d'):
      line=line.split()
      if int(line[1])==1:
        OUTPUT['Property2d']=1
    elif line.startswith('n_property1d'):
      line = line.split()
      OUTPUT['n_property1d']=int(line[1])
    elif line.startswith('n_property2d'):
      line = line.split()
      OUTPUT['n_property2d']=int(line[1])
    elif line.startswith( 'ezero' ):
      line = line.split()
      ezero = float(line[1])
      OUTPUT['ezero']= ezero
    else:
      if 'End of header' in line:
        break
      else:
        continue

  return header, OUTPUT

#==============================================================GET GEOMETRY FROM QM.in FILE===================================================================#
def get_geom(natoms,QMin,Properties):
  #get information of the geometry from the interpolated and aligned xyz-files
  data=open(QMin, 'r').readlines()
  iline=0
  """string_geom='! 11 Geometry in a.u.\n'
  for line in data:
    line.strip()
    iline+=1
    if iline==len(data):
      break
    if line.startswith('tmp.pdb'):
      for atom in range(natoms):
        iline+=1
        line = data[iline]
        line = line.split()
        string_geom+='%20.12f %20.12f %20.12f \n' %(float(line[1]),float(line[2]),float(line[3]))"""
  Geometries=np.zeros((natoms,3))
  for line in data:
    line.strip()
    iline+=1
    if iline==len(data):
      break
    for atom in range(natoms):
      iline+=1
      line=data[iline]
      #print line
      line=line.split()
      #convert from Angstrom to Bohr
      Geometries[atom][0]=float(line[1])/0.529177
      Geometries[atom][1]=float(line[2])/0.529177
      Geometries[atom][2]=float(line[3])/0.529177
    break
  Properties['Geometries'] = Geometries
  return Properties

#============================================================================================================================================================#
#======================================================================PHASECORRECTION=======================================================================#
#============================================================================================================================================================#
def phasecorrection(Properties,oldfilename,filename,newfilename,outputdatfile,inputfile,QMin):
  #corrects the phases of the Hamiltonian (SOCs), NACs and Dipole Moments (and additionally of Overlaps)
  header,OUTPUT=get_header(outputdatfile)
  nmstates=OUTPUT['nmstates']
  natoms = OUTPUT['natoms']
  n_singlets = OUTPUT['n_singlets']
  n_triplets = OUTPUT['n_triplets']
  #print(n_singlets,n_triplets)
  #if n_triplets==None:
  #  n_triplets = int(0)
  Properties,data=read_QMout(Properties,nmstates,n_singlets,n_triplets,filename)
  #get ezero
  Properties=read_input(Properties,inputfile)
  #multiply the phasevector with the previous phasevector
  phasevector_old = Properties['phasevector_original']
  phasevector = Properties['phasevector']
  phasevector = phasevector_old*phasevector
  #print phasevector_old
  #print Properties['phasevector']
  #print phasevector
  Properties.update({'phasevector': phasevector})
  phasevector2=np.zeros((nmstates*2))
  #double each entry of the phasevector to make calculation of real and imaginary values easier
  for numberofstates in range(nmstates):
    phasevector2[numberofstates*2]=phasevector[numberofstates]
    phasevector2[numberofstates*2+1]=phasevector[numberofstates]

  #DO PHASECORRECTION
  QMout=open(newfilename, "w")
  iline=-1
  line_index=-1
  for line in data:
    line = line.strip()
    iline+=1
    line_index+=1
    if iline==len(data):
      break

    elif line.startswith('! 1 Hamiltonian Matrix'):
      #write the next two lines to file
      jline=line_index
      for i in range(2):
        line = data[jline]
        QMout.write(line)
        jline+=1
      #instead of writing the original Hamiltonian, do phasecorrection and write the new Hamiltonian
      Hamiltonian=Properties['Hamiltonian']
      #do calculation of rows with vector
      Hamiltonian_phasecorrected = Hamiltonian * phasevector2
      #do calculation of columns with vector by calculation of each row with the first entry of the vector
      for column in range(nmstates):
        for row_element in range(nmstates*2):
          Hamiltonian_phasecorrected[column][row_element]=Hamiltonian_phasecorrected[column][row_element]*phasevector[column]
      Properties.update({'Hamiltonian_phasecorrected': Hamiltonian_phasecorrected})
      #write the Hamiltonian to QM.out (next 8 lines)
      string_hamiltonian=''
      for hamilton in range(nmstates):
        for row_elements in range(nmstates*2):
          string_hamiltonian+='%20.12f '%Hamiltonian_phasecorrected[hamilton][row_elements]
        string_hamiltonian+='\n'
      string_hamiltonian+='\n'
      QMout.write(string_hamiltonian)
      #substract ezero from energy values of Hamiltonian for output.dat
      ezero = Properties['ezero']
      Hamiltonian_phasecorrected_ezero=Hamiltonian_phasecorrected
      for numberofstates in range(nmstates):
        for energyvalues in range(nmstates):
          if numberofstates == energyvalues:
            Hamiltonian_phasecorrected_ezero[numberofstates][energyvalues*2]=Hamiltonian_phasecorrected[numberofstates][energyvalues*2]-ezero
      Properties.update({'Hamiltonian_phasecorrected_ezero': Hamiltonian_phasecorrected_ezero})

    elif line.startswith('! 2 Dipole Moment'):
      #write the next line to the file
      jline = line_index
      line=data[jline]
      QMout.write(line)
      #do phasecorrection
      Dipole_x=Properties['Dipole_x']
      Dipole_y=Properties['Dipole_y']
      Dipole_z=Properties['Dipole_z']
      #
      Dipole_x_phasecorrected = []
      Dipole_y_phasecorrected = []
      Dipole_z_phasecorrected = []
      #
      Dipole_x_phasecorrected = Dipole_x * phasevector2
      Dipole_y_phasecorrected = Dipole_y * phasevector2
      Dipole_z_phasecorrected = Dipole_z * phasevector2
      #
      for column in range(nmstates):
        for row_element in range(nmstates*2):
          Dipole_x_phasecorrected[column][row_element]=Dipole_x_phasecorrected[column][row_element]*phasevector[column]
          Dipole_y_phasecorrected[column][row_element]=Dipole_y_phasecorrected[column][row_element]*phasevector[column]
          Dipole_z_phasecorrected[column][row_element]=Dipole_z_phasecorrected[column][row_element]*phasevector[column]
      #
      Properties.update({'Dipole_x_phasecorrected': Dipole_x_phasecorrected})
      Properties.update({'Dipole_y_phasecorrected': Dipole_y_phasecorrected})
      Properties.update({'Dipole_z_phasecorrected': Dipole_z_phasecorrected})
      #
      string_dipole=''
      jline+=1
      line=data[jline]
      string_dipole+='%s' %line
      for dipole in range(nmstates):
        for row_elements in range(nmstates*2):
          string_dipole+='%20.12f '%Dipole_x_phasecorrected[dipole][row_elements]
        string_dipole+='\n'
        jline+=1
      jline+=1
      line = data[jline]
      string_dipole+='%s' %line
      for dipole in range(nmstates):
        for row_elements in range(nmstates*2):
          string_dipole+='%20.12f '%Dipole_y_phasecorrected[dipole][row_elements]
        string_dipole+='\n'
        jline+=1
      jline+=1
      line = data[jline]
      string_dipole+='%s'%line
      for dipole in range(nmstates):
        for row_elements in range(nmstates*2):
          string_dipole+='%20.12f '%Dipole_z_phasecorrected[dipole][row_elements]
        string_dipole+='\n'
        jline+=1
      #print string_dipole
      QMout.write(string_dipole)

    elif line.startswith('! 3 Gradient Vectors'):
      #only write lines - no phasecorrection for gradients
      #write the next lines containing gradient information to the file to the file
      jline = line_index
      for gradient_index in range((natoms+1)*nmstates+1):
        line=data[jline]
        QMout.write(line)
        jline+=1

    elif line.startswith('! 5 Non-adiabatic couplings'):
      #write first line to the QM.out file
      jline = line_index
      line = data[jline]
      QMout.write(line)

      #do phasecorrection
      NAC=Properties['NonAdiabaticCouplings']
      #the NAC matrix is saved as a matrix containing natoms*3 entries per line (so one line contains the nacs between a specific state and another specific one)
      #each line will be corrected with the phase - first, every line has to be corrected with the first entry of the phasecorrection vector, the second line with the second,...
      NAC_phasecorrected = np.zeros((nmstates*nmstates,natoms*3))

      #correct by multiplication with phasecorrection_vector line by line
      phasecorrection_index_line=0
      for nacindex_line in range(nmstates*nmstates):
        for nac_state_line in range(natoms*3):
          NAC_phasecorrected[nacindex_line][nac_state_line] = NAC[nacindex_line][nac_state_line] * phasevector[phasecorrection_index_line]
        phasecorrection_index_line+=1
        if phasecorrection_index_line >= nmstates:
          phasecorrection_index_line=0

      #do phasecorrection column by column
      #this means, the first nmstates*line entries have to be multiplicated by the first entry of the phasecorrection vector
      phasecorrection_index_column=0
      iterator=0
      for nacindex_column in range(nmstates*nmstates):
        for nac_state_column in range(natoms*3):
          NAC_phasecorrected[nacindex_column][nac_state_column] = NAC_phasecorrected[nacindex_column][nac_state_column] * phasevector[phasecorrection_index_column]
        iterator+=1
        if iterator >= nmstates:
          iterator=0
          phasecorrection_index_column+=1
      Properties.update({'NAC_phasecorrected': NAC_phasecorrected})

      #make string for writing to QM.out file
      string_nac=''
      for nacstring in range(nmstates*nmstates):
        #contains information about the states of the NACs
        jline+=1
        line=data[jline]
        string_nac+='%s' %line
        jline+=natoms
        xyz_index=0
        for nacline in range(natoms*3):
          xyz_index+=1
          string_nac+='%20.12f ' %NAC_phasecorrected[nacstring][nacline]
          if xyz_index>=3:
            xyz_index=0
            string_nac+='\n'
      QMout.write(string_nac)

    elif line.startswith('! 6 Overlap matrix'):
      #write the next 2 lines to the QM.out file
      jline = line_index
      for index in range(2):
        line = data[jline]
        QMout.write(line)
        jline+=1

      Overlap=Properties['Overlap']
      #get the old phasevector since the overlaps are mutiplied line by line with the current phasevector and column by column with the one of the previous timestep (<S1|S2>)
      phasevector_old=Properties['phasevector_original']
      #first phasecorrection line by line
      Overlap_phasecorrected = Overlap * phasevector2
      #second phasecorrection column by column
      for lineindex in range(nmstates):
        for overlap_index in range(nmstates*2):
          Overlap_phasecorrected[lineindex][overlap_index] = Overlap_phasecorrected[lineindex][overlap_index] * phasevector_old[lineindex]
      #prepare String
      string_overlap=''
      for overlap in range(nmstates):
        for line_value in range(nmstates*2):
          string_overlap+='%20.12f ' %Overlap_phasecorrected[overlap][line_value]
        string_overlap+='\n'
      string_overlap+='\n'
      QMout.write(string_overlap)

    elif line.startswith('! 7 Phases'):
      #writes the phase vector
      jline = line_index
      for phasestring in range(2):
        line = data[jline]
        QMout.write(line)
        jline+=1
      string_phase=''
      for phases in range(nmstates):
        string_phase+='%20.12f 0.000000000\n' %phasevector[phases]
      QMout.write(string_phase)

    else:
      pass

  QMout.close()

  Properties=get_geom(natoms,QMin,Properties)

  return Properties, OUTPUT

#======================================================================READ QM.out===================================================================#
#complementary to function "check_phases" from "NN_phasecorrection.py"

def read_QMinit_out(oldfilename,QMin,OUTPUT):
  Properties={}
  #get the phasevector
  readQMout=open(oldfilename,'r')
  data=readQMout.readlines()
  iline=-1
  for line in data:
    iline+=1
    line = line.strip()
    if iline==len(data):
      break
    line=data[1]
    t = line.split()
    nmstates=t[0]
    nmstates=int(nmstates)
    Properties['nmstates']=nmstates
    break
  iline = -1
  for line in data:
    iline += 1
    line = data[iline]
    line = line.strip()
    if iline == len(data):
      break
    elif line.startswith('! 7 Phases'):
      iline+=1
      line=data[iline]
      line=line.split()
      phasevector=np.zeros((nmstates))
      for numberofstates in range(nmstates):
        line = data[iline+1+numberofstates]
        line = line.split()
        #print line
        phasevector[numberofstates]=line[0]
      Properties['phasevector_original']=phasevector
      #print Properties['phasevector_original']
      break
    else:
      #set the phasevector to +1
      phasevector=np.ones((nmstates))
      Properties['phasevector_original']=phasevector
  Properties,data=read_QMout(Properties,nmstates,OUTPUT['n_singlets'], OUTPUT['n_triplets'],oldfilename)
  Properties=get_geom(OUTPUT['natoms'],QMin,Properties)
  return Properties

def read_QMout(Properties,nmstates,n_singlets,n_triplets,filename):
  #iterate over QM.out files
  Properties.update({'phasevector':		False,
			  'Hamiltonian':		False,
			  'Dipole_x':   		False,
			  'Dipole_y':   		False,
			  'Dipole_z':   		False,
			  'Gradient':  			False,
			  'NonAdibaticCouplings':	False,
			  'Overlap':			False})
  #start from 1 because the inital file does not have to be included
  readQMout=open(filename, "r")
  data=readQMout.readlines()
  readQMout.close()
  overlap_vector_correct=np.zeros((nmstates))
  threshold = float(0.5)
  iline=-1
  #print( n_singlets, n_triplets)
  for line in data:
    line=line.strip()
    iline+=1
    if line.startswith("! 6 Overlap matrix" ):
      iline+=1
      #check phases of singlets only
      for overlaps_singlets in range(n_singlets):
        iline+=1
        line = data[iline]
        line = line.split()
        overlap_singlets = float(line[2*overlaps_singlets])
        #print( overlap)
        if abs(overlap_singlets) < threshold:
          for overlaps_stateswitch in range(n_singlets):
            overlap_switched_state = float(line[2*overlaps_stateswitch])
            #print (overlap_switched_state)
            if abs(overlap_switched_state) > threshold:
              #print("states switched")
              overlap_singlets = overlap_switched_state
              if overlap_singlets < int(0):
                overlap_singlets = int(-1)
              else:
                overlap_singlets = int(1)
        else:
          if overlap_singlets < int(0):
            overlap_singlets = int(-1)
          else:
            overlap_singlets = int(1)
        overlap_vector_correct[overlaps_singlets]=overlap_singlets
      #check phases of triplets - do this three times since triplets are three times degenerated
      triplet_nmstate_0=int(0)
      triplet_nmstate_1=int(0)
      if n_triplets==0:
        pass
      else:
        for triplet_number in range(1):
          triplet_nmstate_0=triplet_nmstate_1
          triplet_nmstate_1=triplet_nmstate_0+1
          #print(triplet_magnetic, "should start with 1")
          for overlaps_triplets in range(n_singlets*triplet_nmstate_0,n_singlets+triplet_nmstate_1*n_triplets):
            iline+=1
            line = data[iline]
            line = line.split()
            overlap_triplets = float(line[2*overlaps_triplets])
            if abs(overlap_triplets) < threshold:
              for overlaps_stateswitch_triplets in range(n_singlets*triplet_nmstate_0,n_singlets+triplet_nmstate_1*n_triplets):
                overlap_switched_state_triplet = float(line[2*overlaps_stateswitch_triplets])
                if abs(overlap_switched_state_triplet) > threshold:
                  overlap_triplets = overlap_switched_state_triplet
                  if overlap_triplets < int(0):
                    overlap_triplets = int(-1)
                  else:
                    overlap_triplets = int(1)
            else:
              if overlap_triplets < int(0):
                overlap_triplets = int(-1)
              else:
                overlap_triplets = int(1)
            overlap_vector_correct[overlaps_triplets]=overlap_triplets
    #get number of states and write out vector with phases
    """if line.startswith("! 7 Phases"):
      phasevector=np.zeros((nmstates))
      for numberofstates in range(nmstates):
        line=data[iline+2+numberofstates]
        line=line.split()
        phasevector[numberofstates]=line[0]
      Properties['phasevector']=phasevector
      break"""
  iline=-1
  #print(overlap_vector_correct)
  Properties['phasevector']=overlap_vector_correct
  for line in data:
    line=line.strip()
    iline+=1
    if line.startswith('! 1 Hamiltonian Matrix'):
      Hamiltonian=np.zeros((nmstates,nmstates*2))
      for numberofstates in range(nmstates):
        line=data[iline+2+numberofstates]
        line=line.split()
        for real_imag in range(nmstates*2):
          Hamiltonian[numberofstates][real_imag]=line[real_imag]
      Properties['Hamiltonian']=Hamiltonian
    elif line.startswith('! 2 Dipole Moment Matrices'):
      Dipole_x=np.zeros((nmstates,nmstates*2))
      Dipole_y=np.zeros((nmstates,nmstates*2))
      Dipole_z=np.zeros((nmstates,nmstates*2))
      for numberofstates in range(nmstates):
        line_x=data[iline+2+numberofstates]
        line_x=line_x.split()
        line_y=data[iline+3+nmstates+numberofstates]
        line_y=line_y.split()
        line_z=data[iline+4+nmstates*2+numberofstates]
        line_z=line_z.split()
        for real_imag in range(nmstates*2):
          Dipole_x[numberofstates][real_imag]=line_x[real_imag]
          Dipole_y[numberofstates][real_imag]=line_y[real_imag]
          Dipole_z[numberofstates][real_imag]=line_z[real_imag]
      Properties['Dipole_x']=Dipole_x
      Properties['Dipole_y']=Dipole_y
      Properties['Dipole_z']=Dipole_z
    elif line.startswith('! 3 Gradient Vectors'):
      line=data[iline+1]
      line=line.split()
      natoms=int(line[0])
      Properties.update({'natoms': natoms})
      #writes the gradient tensor in form of a matrix  -  every natomsx(x,y,z) matrix is written as one vector
      Gradient=np.zeros((nmstates,natoms*3))
      index=0
      for numberofstates in range(nmstates):
        index+=1
        gradientindex=0
        for atom_index in range(natoms):
          index+=1
          line=data[iline+index]
          line=line.split()
          for xyz in range(3):
            Gradient[numberofstates][gradientindex]=line[xyz]
            gradientindex+=1
      Properties['Gradient']=Gradient
    elif line.startswith('! 5 Non-adiabatic couplings'):
      #the nacs will be written as a matrix with the size: (numberofstatesxnumberofstates)x(natoms*3(xyz))
      NAC=np.zeros((nmstates*nmstates,natoms*3))
      index=0
      for numberofstates in range(nmstates*nmstates):
        index+=1
        nacindex=0
        for atom_index in range(natoms):
          index+=1
          line=data[iline+index]
          line=line.split()
          for xyz in range(3):
            NAC[numberofstates][nacindex]=line[xyz]
            nacindex+=1
      Properties['NonAdiabaticCouplings']=NAC
    elif line.startswith('! 6 Overlap matrix'):
      overlap=np.zeros((nmstates,nmstates*2))
      for numberofstates in range(nmstates):
        line=data[iline+2+numberofstates]
        line=line.split()
        for real_imag in range(nmstates*2):
          overlap[numberofstates][real_imag]=line[real_imag]
      Properties['Overlap']=overlap

  return Properties, data

def read_input(Properties, inputfile):
  inputf=open(inputfile, 'r').readlines()
  for line in inputf:
    line=line.strip()
    if line.startswith( 'ezero' ):
      line = line.split()
      ezero = float(line[1])
      Properties.update({'ezero': ezero})
  return Properties

if __name__ == "__main__":
  try:
    name,outputdatforheader, QMin = argv
  except ValueError:
    print( "Usage: script <output.dat file for header information> <file of geometry used for calcualtion of QMout; often QM.in>")
    exit()
  outputdatfile = argv[1]
  filename=str("QM.out")
  inputfile=argv[2]
  header,OUTPUT=get_header(outputdatfile)
  Properties=read_QMinit_out(filename,QMin,OUTPUT)
  #Properties,OUTPUT=phasecorrection(Properties,filename,outputdatfile,inputfile)
  qm2outputdat(Properties,outputdatfile)


"""
Class interface to the model calculator.

Calling a model is somewhat non-trivial since the functions called depend
on the data type.  For 1D data the *Iq* kernel needs to be called, for
2D data the *Iqxy* kernel needs to be called, and for SESANS data the
*Iq* kernel needs to be called followed by a Hankel transform.  Before
the kernel is called an appropriate *q* calculation vector needs to be
constructed.  This is not the simple *q* vector where you have measured
the data since the resolution calculation will require values beyond the
range of the measured data.  After the calculation the resolution calculator
must be called to return the predicted value for each measured data point.

:class:`DirectModel` is a callable object that takes *parameter=value*
keyword arguments and returns the appropriate theory values for the data.

:class:`DataMixin` does the real work of interpreting the data and calling
the model calculator.  This is used by :class:`DirectModel`, which uses
direct parameter values and by :class:`bumps_model.Experiment` which wraps
the parameter values in boxes so that the user can set fitting ranges, etc.
on the individual parameters and send the model to the Bumps optimizers.
"""
from __future__ import print_function

import numpy as np  # type: ignore

# TODO: fix sesans module
from . import sesans  # type: ignore
from . import weights
from . import resolution
from . import resolution2d
from .details import make_kernel_args, dispersion_mesh

try:
    from typing import Optional, Dict, Tuple
except ImportError:
    pass
else:
    from .data import Data
    from .kernel import Kernel, KernelModel
    from .modelinfo import Parameter, ParameterSet

def call_kernel(calculator, pars, cutoff=0., mono=False):
    # type: (Kernel, ParameterSet, float, bool) -> np.ndarray
    """
    Call *kernel* returned from *model.make_kernel* with parameters *pars*.

    *cutoff* is the limiting value for the product of dispersion weights used
    to perform the multidimensional dispersion calculation more quickly at a
    slight cost to accuracy. The default value of *cutoff=0* integrates over
    the entire dispersion cube.  Using *cutoff=1e-5* can be 50% faster, but
    with an error of about 1%, which is usually less than the measurement
    uncertainty.

    *mono* is True if polydispersity should be set to none on all parameters.
    """
    parameters = calculator.info.parameters
    if mono:
        active = lambda name: False
    elif calculator.dim == '1d':
        active = lambda name: name in parameters.pd_1d
    elif calculator.dim == '2d':
        active = lambda name: name in parameters.pd_2d
    else:
        active = lambda name: True

    #print("pars",[p.id for p in parameters.call_parameters])
    vw_pairs = [(get_weights(p, pars) if active(p.name)
                 else ([pars.get(p.name, p.default)], [1.0]))
                for p in parameters.call_parameters]

    call_details, values, is_magnetic = make_kernel_args(calculator, vw_pairs)
    #print("values:", values)
    return calculator(call_details, values, cutoff, is_magnetic)

def call_ER(model_info, pars):
    # type: (ModelInfo, ParameterSet) -> float
    """
    Call the model ER function using *values*.

    *model_info* is either *model.info* if you have a loaded model,
    or *kernel.info* if you have a model kernel prepared for evaluation.
    """
    if model_info.ER is None:
        return 1.0
    elif not model_info.parameters.form_volume_parameters:
        # handle the case where ER is provided but model is not polydisperse
        return model_info.ER()
    else:
        value, weight = _vol_pars(model_info, pars)
        individual_radii = model_info.ER(*value)
        return np.sum(weight*individual_radii) / np.sum(weight)

def call_VR(model_info, pars):
    # type: (ModelInfo, ParameterSet) -> float
    """
    Call the model VR function using *pars*.

    *model_info* is either *model.info* if you have a loaded model,
    or *kernel.info* if you have a model kernel prepared for evaluation.
    """
    if model_info.VR is None:
        return 1.0
    elif not model_info.parameters.form_volume_parameters:
        # handle the case where ER is provided but model is not polydisperse
        return model_info.VR()
    else:
        value, weight = _vol_pars(model_info, pars)
        whole, part = model_info.VR(*value)
        return np.sum(weight*part)/np.sum(weight*whole)

def call_profile(model_info, **pars):
    # type: (ModelInfo, ...) -> Tuple[np.ndarray, np.ndarray, Tuple[str, str]]
    """
    Returns the profile *x, y, (xlabel, ylabel)* representing the model.
    """
    args = {}
    for p in model_info.parameters.kernel_parameters:
        if p.length > 1:
            value = np.array([pars.get(p.id+str(j), p.default)
                              for j in range(1, p.length+1)])
        else:
            value = pars.get(p.id, p.default)
        args[p.id] = value
    x, y = model_info.profile(**args)
    return x, y, model_info.profile_axes

def get_weights(parameter, values):
    # type: (Parameter, Dict[str, float]) -> Tuple[np.ndarray, np.ndarray]
    """
    Generate the distribution for parameter *name* given the parameter values
    in *pars*.

    Uses "name", "name_pd", "name_pd_type", "name_pd_n", "name_pd_sigma"
    from the *pars* dictionary for parameter value and parameter dispersion.
    """
    value = float(values.get(parameter.name, parameter.default))
    relative = parameter.relative_pd
    limits = parameter.limits
    disperser = values.get(parameter.name+'_pd_type', 'gaussian')
    npts = values.get(parameter.name+'_pd_n', 0)
    width = values.get(parameter.name+'_pd', 0.0)
    nsigma = values.get(parameter.name+'_pd_nsigma', 3.0)
    if npts == 0 or width == 0:
        return [value], [1.0]
    value, weight = weights.get_weights(
        disperser, npts, width, nsigma, value, limits, relative)
    return value, weight / np.sum(weight)

def _vol_pars(model_info, pars):
    # type: (ModelInfo, ParameterSet) -> Tuple[np.ndarray, np.ndarray]
    vol_pars = [get_weights(p, pars)
                for p in model_info.parameters.call_parameters
                if p.type == 'volume']
    #import pylab; pylab.plot(vol_pars[0][0],vol_pars[0][1]); pylab.show()
    value, weight = dispersion_mesh(model_info, vol_pars)
    return value, weight

class DataMixin(object):
    """
    DataMixin captures the common aspects of evaluating a SAS model for a
    particular data set, including calculating Iq and evaluating the
    resolution function.  It is used in particular by :class:`DirectModel`,
    which evaluates a SAS model parameters as key word arguments to the
    calculator method, and by :class:`bumps_model.Experiment`, which wraps the
    model and data for use with the Bumps fitting engine.  It is not
    currently used by :class:`sasview_model.SasviewModel` since this will
    require a number of changes to SasView before we can do it.

    :meth:`_interpret_data` initializes the data structures necessary
    to manage the calculations.  This sets attributes in the child class
    such as *data_type* and *resolution*.

    :meth:`_calc_theory` evaluates the model at the given control values.

    :meth:`_set_data` sets the intensity data in the data object,
    possibly with random noise added.  This is useful for simulating a
    dataset with the results from :meth:`_calc_theory`.
    """
    def _interpret_data(self, data, model):
        # type: (Data, KernelModel) -> None
        # pylint: disable=attribute-defined-outside-init

        self._data = data
        self._model = model

        # interpret data
        if hasattr(data, 'isSesans') and data.isSesans:
            self.data_type = 'sesans'
        elif hasattr(data, 'qx_data'):
            self.data_type = 'Iqxy'
        elif getattr(data, 'oriented', False):
            self.data_type = 'Iq-oriented'
        else:
            self.data_type = 'Iq'

        if self.data_type == 'sesans':
            q = sesans.make_q(data.sample.zacceptance, data.Rmax)
            index = slice(None, None)
            res = None
            if data.y is not None:
                Iq, dIq = data.y, data.dy
            else:
                Iq, dIq = None, None
            #self._theory = np.zeros_like(q)
            q_vectors = [q]
            q_mono = sesans.make_all_q(data)
        elif self.data_type == 'Iqxy':
            #if not model.info.parameters.has_2d:
            #    raise ValueError("not 2D without orientation or magnetic parameters")
            q = np.sqrt(data.qx_data**2 + data.qy_data**2)
            qmin = getattr(data, 'qmin', 1e-16)
            qmax = getattr(data, 'qmax', np.inf)
            accuracy = getattr(data, 'accuracy', 'Low')
            index = ~data.mask & (q >= qmin) & (q <= qmax)
            if data.data is not None:
                index &= ~np.isnan(data.data)
                Iq = data.data[index]
                dIq = data.err_data[index]
            else:
                Iq, dIq = None, None
            res = resolution2d.Pinhole2D(data=data, index=index,
                                         nsigma=3.0, accuracy=accuracy)
            #self._theory = np.zeros_like(self.Iq)
            q_vectors = res.q_calc
            q_mono = []
        elif self.data_type == 'Iq':
            index = (data.x >= data.qmin) & (data.x <= data.qmax)
            if data.y is not None:
                index &= ~np.isnan(data.y)
                Iq = data.y[index]
                dIq = data.dy[index]
            else:
                Iq, dIq = None, None
            if getattr(data, 'dx', None) is not None:
                q, dq = data.x[index], data.dx[index]
                if (dq > 0).any():
                    res = resolution.Pinhole1D(q, dq)
                else:
                    res = resolution.Perfect1D(q)
            elif (getattr(data, 'dxl', None) is not None
                  and getattr(data, 'dxw', None) is not None):
                res = resolution.Slit1D(data.x[index],
                                        qx_width=data.dxl[index],
                                        qy_width=data.dxw[index])
            else:
                res = resolution.Perfect1D(data.x[index])

            #self._theory = np.zeros_like(self.Iq)
            q_vectors = [res.q_calc]
            q_mono = []
        elif self.data_type == 'Iq-oriented':
            index = (data.x >= data.qmin) & (data.x <= data.qmax)
            if data.y is not None:
                index &= ~np.isnan(data.y)
                Iq = data.y[index]
                dIq = data.dy[index]
            else:
                Iq, dIq = None, None
            if (getattr(data, 'dxl', None) is None
                    or getattr(data, 'dxw', None) is None):
                raise ValueError("oriented sample with 1D data needs slit resolution")

            res = resolution2d.Slit2D(data.x[index],
                                      qx_width=data.dxw[index],
                                      qy_width=data.dxl[index])
            q_vectors = res.q_calc
            q_mono = []
        else:
            raise ValueError("Unknown data type") # never gets here

        # Remember function inputs so we can delay loading the function and
        # so we can save/restore state
        self._kernel_inputs = q_vectors
        self._kernel_mono_inputs = q_mono
        self._kernel = None
        self.Iq, self.dIq, self.index = Iq, dIq, index
        self.resolution = res

    def _set_data(self, Iq, noise=None):
        # type: (np.ndarray, Optional[float]) -> None
        # pylint: disable=attribute-defined-outside-init
        if noise is not None:
            self.dIq = Iq*noise*0.01
        dy = self.dIq
        y = Iq + np.random.randn(*dy.shape) * dy
        self.Iq = y
        if self.data_type in ('Iq', 'Iq-oriented'):
            self._data.dy[self.index] = dy
            self._data.y[self.index] = y
        elif self.data_type == 'Iqxy':
            self._data.data[self.index] = y
        elif self.data_type == 'sesans':
            self._data.y[self.index] = y
        else:
            raise ValueError("Unknown model")

    def _calc_theory(self, pars, cutoff=0.0):
        # type: (ParameterSet, float) -> np.ndarray
        if self._kernel is None:
            self._kernel = self._model.make_kernel(self._kernel_inputs)
            self._kernel_mono = (
                self._model.make_kernel(self._kernel_mono_inputs)
                if self._kernel_mono_inputs else None)

        Iq_calc = call_kernel(self._kernel, pars, cutoff=cutoff)
        # Storing the calculated Iq values so that they can be plotted.
        # Only applies to oriented USANS data for now.
        # TODO: extend plotting of calculate Iq to other measurement types
        # TODO: refactor so we don't store the result in the model
        self.Iq_calc = None
        if self.data_type == 'sesans':
            Iq_mono = (call_kernel(self._kernel_mono, pars, mono=True)
                       if self._kernel_mono_inputs else None)
            result = sesans.transform(self._data,
                                      self._kernel_inputs[0], Iq_calc,
                                      self._kernel_mono_inputs, Iq_mono)
        else:
            result = self.resolution.apply(Iq_calc)
            if hasattr(self.resolution, 'nx'):
                self.Iq_calc = (
                    self.resolution.qx_calc, self.resolution.qy_calc,
                    np.reshape(Iq_calc, (self.resolution.ny, self.resolution.nx))
                )
        return result

class DirectModel(DataMixin):
    """
    Create a calculator object for a model.

    *data* is 1D SAS, 2D SAS or SESANS data

    *model* is a model calculator return from :func:`generate.load_model`

    *cutoff* is the polydispersity weight cutoff.
    """
    def __init__(self, data, model, cutoff=1e-5):
        # type: (Data, KernelModel, float) -> None
        self.model = model
        self.cutoff = cutoff
        # Note: _interpret_data defines the model attributes
        self._interpret_data(data, model)

    def __call__(self, **pars):
        # type: (**float) -> np.ndarray
        return self._calc_theory(pars, cutoff=self.cutoff)

    def simulate_data(self, noise=None, **pars):
        # type: (Optional[float], **float) -> None
        """
        Generate simulated data for the model.
        """
        Iq = self.__call__(**pars)
        self._set_data(Iq, noise=noise)

    def profile(self, **pars):
        # type: (**float) -> None
        """
        Generate a plottable profile.
        """
        return call_profile(self.model.info, **pars)

def main():
    # type: () -> None
    """
    Program to evaluate a particular model at a set of q values.
    """
    import sys
    from .data import empty_data1D, empty_data2D
    from .core import load_model_info, build_model

    if len(sys.argv) < 3:
        print("usage: python -m sasmodels.direct_model modelname (q|qx,qy) par=val ...")
        sys.exit(1)
    model_name = sys.argv[1]
    call = sys.argv[2].upper()
    if call != "ER_VR":
        try:
            values = [float(v) for v in call.split(',')]
        except Exception:
            values = []
        if len(values) == 1:
            q, = values
            data = empty_data1D([q])
        elif len(values) == 2:
            qx, qy = values
            data = empty_data2D([qx], [qy])
        else:
            print("use q or qx,qy or ER or VR")
            sys.exit(1)
    else:
        data = empty_data1D([0.001])  # Data not used in ER/VR

    model_info = load_model_info(model_name)
    model = build_model(model_info)
    calculator = DirectModel(data, model)
    pars = dict((k, (float(v) if not k.endswith("_pd_type") else v))
                for pair in sys.argv[3:]
                for k, v in [pair.split('=')])
    if call == "ER_VR":
        ER = call_ER(model_info, pars)
        VR = call_VR(model_info, pars)
        print(ER, VR)
    else:
        Iq = calculator(**pars)
        print(Iq[0])

if __name__ == "__main__":
    main()


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

# This code is part of Qiskit.
#
# (C) Copyright IBM 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Parametric waveforms module. These are pulses which are described by a specified
parameterization.

If a backend supports parametric pulses, it will have the attribute
`backend.configuration().parametric_pulses`, which is a list of supported pulse shapes, such as
`['gaussian', 'gaussian_square', 'drag']`. A Pulse Schedule, using parametric pulses, which is
assembled for a backend which supports those pulses, will result in a Qobj which is dramatically
smaller than one which uses Waveforms.

This module can easily be extended to describe more pulse shapes. The new class should:
  - have a descriptive name
  - be a well known and/or well described formula (include the formula in the class docstring)
  - take some parameters (at least `duration`) and validate them, if necessary
  - implement a `get_waveform` method which returns a corresponding Waveform in the
    case that it is assembled for a backend which does not support it.

The new pulse must then be registered by the assembler in
`qiskit/qobj/converters/pulse_instruction.py:ParametricPulseShapes`
by following the existing pattern:

    class ParametricPulseShapes(Enum):
        gaussian = pulse_lib.Gaussian
        ...
        new_supported_pulse_name = pulse_lib.YourPulseWaveformClass
"""
import warnings
from abc import abstractmethod
from typing import Any, Callable, Dict, Optional
import math
import numpy as np

from . import continuous
from .discrete import gaussian, gaussian_square, drag, constant
from .pulse import Pulse
from .waveform import Waveform
from ..exceptions import PulseError

class ParametricPulse(Pulse):
    """The abstract superclass for parametric pulses."""

    @abstractmethod
    def __init__(self, duration: int, name: Optional[str] = None):
        """Create a parametric pulse and validate the input parameters.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            name: Display name for this pulse envelope.
        """
        super().__init__(duration=duration, name=name)
        self.validate_parameters()

    @abstractmethod
    def get_waveform(self) -> Waveform:
        """Return a Waveform with samples filled according to the formula that the pulse
        represents and the parameter values it contains.
        """
        raise NotImplementedError

    def get_sample_pulse(self) -> Waveform:
        """Deprecated."""
        warnings.warn('`get_sample_pulse` has been deprecated. '
                      ' Use `get_waveform` instead.', DeprecationWarning)
        return self.get_waveform()

    @abstractmethod
    def validate_parameters(self) -> None:
        """
        Validate parameters.

        Raises:
            PulseError: If the parameters passed are not valid.
        """
        raise NotImplementedError

    @property
    @abstractmethod
    def parameters(self) -> Dict[str, Any]:
        """Return a dictionary containing the pulse's parameters."""
        pass

    def draw(self, dt: float = 1,
             style=None,
             filename: Optional[str] = None,
             interp_method: Optional[Callable] = None,
             scale: float = 1, interactive: bool = False):
        """Plot the pulse.

        Args:
            dt: Time interval of samples.
            style (Optional[PulseStyle]): A style sheet to configure plot appearance
            filename: Name required to save pulse image
            interp_method: A function for interpolation
            scale: Relative visual scaling of waveform amplitudes
            interactive: When set true show the circuit in a new window
                (this depends on the matplotlib backend being used supporting this)

        Returns:
            matplotlib.figure: A matplotlib figure object of the pulse envelope
        """
        return self.get_waveform().draw(dt=dt, style=style, filename=filename,
                                        interp_method=interp_method, scale=scale,
                                        interactive=interactive)

    def __eq__(self, other: Pulse) -> bool:
        return super().__eq__(other) and self.parameters == other.parameters

    def __hash__(self) -> int:
        return hash(self.parameters[k] for k in sorted(self.parameters))

class Gaussian(ParametricPulse):
    """A truncated pulse envelope shaped according to the Gaussian function whose mean is centered
    at the center of the pulse (duration / 2):

    .. math::

        f(x) = amp * exp( -(1/2) * (x - duration/2)^2 / sigma^2) )  ,  0 <= x < duration
    """

    def __init__(self,
                 duration: int,
                 amp: complex,
                 sigma: float,
                 name: Optional[str] = None):
        """Initialize the gaussian pulse.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            amp: The amplitude of the Gaussian envelope.
            sigma: A measure of how wide or narrow the Gaussian peak is; described mathematically
                   in the class docstring.
            name: Display name for this pulse envelope.
        """
        self._amp = complex(amp)
        self._sigma = sigma
        super().__init__(duration=duration, name=name)

    @property
    def amp(self) -> complex:
        """The Gaussian amplitude."""
        return self._amp

    @property
    def sigma(self) -> float:
        """The Gaussian standard deviation of the pulse width."""
        return self._sigma

    def get_waveform(self) -> Waveform:
        return gaussian(duration=self.duration, amp=self.amp,
                        sigma=self.sigma, zero_ends=False)

    def validate_parameters(self) -> None:
        if abs(self.amp) > 1.:
            raise PulseError("The amplitude norm must be <= 1, "
                             "found: {}".format(abs(self.amp)))
        if self.sigma <= 0:
            raise PulseError("Sigma must be greater than 0.")

    @property
    def parameters(self) -> Dict[str, Any]:
        return {"duration": self.duration, "amp": self.amp, "sigma": self.sigma}

    def __repr__(self) -> str:
        return "{}(duration={}, amp={}, sigma={}{})" \
               "".format(self.__class__.__name__, self.duration, self.amp, self.sigma,
                         ", name='{}'".format(self.name) if self.name is not None else "")

class GaussianSquare(ParametricPulse):
    """A square pulse with a Gaussian shaped risefall on either side:

    .. math::

        risefall = (duration - width) / 2

        0 <= x < risefall

        f(x) = amp * exp( -(1/2) * (x - risefall/2)^2 / sigma^2) )

        risefall <= x < risefall + width

        f(x) = amp

        risefall + width <= x < duration

        f(x) = amp * exp( -(1/2) * (x - (risefall + width)/2)^2 / sigma^2) )
    """

    def __init__(self,
                 duration: int,
                 amp: complex,
                 sigma: float,
                 width: float,
                 name: Optional[str] = None):
        """Initialize the gaussian square pulse.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            amp: The amplitude of the Gaussian and of the square pulse.
            sigma: A measure of how wide or narrow the Gaussian risefall is; see the class
                   docstring for more details.
            width: The duration of the embedded square pulse.
            name: Display name for this pulse envelope.
        """
        self._amp = complex(amp)
        self._sigma = sigma
        self._width = width
        super().__init__(duration=duration, name=name)

    @property
    def amp(self) -> complex:
        """The Gaussian amplitude."""
        return self._amp

    @property
    def sigma(self) -> float:
        """The Gaussian standard deviation of the pulse width."""
        return self._sigma

    @property
    def width(self) -> float:
        """The width of the square portion of the pulse."""
        return self._width

    def get_waveform(self) -> Waveform:
        return gaussian_square(duration=self.duration, amp=self.amp,
                               width=self.width, sigma=self.sigma,
                               zero_ends=False)

    def validate_parameters(self) -> None:
        if abs(self.amp) > 1.:
            raise PulseError("The amplitude norm must be <= 1, "
                             "found: {}".format(abs(self.amp)))
        if self.sigma <= 0:
            raise PulseError("Sigma must be greater than 0.")
        if self.width < 0 or self.width >= self.duration:
            raise PulseError("The pulse width must be at least 0 and less than its duration.")

    @property
    def parameters(self) -> Dict[str, Any]:
        return {"duration": self.duration, "amp": self.amp, "sigma": self.sigma,
                "width": self.width}

    def __repr__(self) -> str:
        return "{}(duration={}, amp={}, sigma={}, width={}{})" \
               "".format(self.__class__.__name__, self.duration, self.amp, self.sigma, self.width,
                         ", name='{}'".format(self.name) if self.name is not None else "")

class Drag(ParametricPulse):
    r"""The Derivative Removal by Adiabatic Gate (DRAG) pulse is a standard Gaussian pulse
    with an additional Gaussian derivative component. It is designed to reduce the frequency
    spectrum of a normal gaussian pulse near the :math:`|1\rangle` - :math:`|2\rangle` transition,
    reducing the chance of leakage to the :math:`|2\rangle` state.

    .. math::

        f(x) = Gaussian + 1j * beta * d/dx [Gaussian]
             = Gaussian + 1j * beta * (-(x - duration/2) / sigma^2) [Gaussian]

    where 'Gaussian' is:

    .. math::

        Gaussian(x, amp, sigma) = amp * exp( -(1/2) * (x - duration/2)^2 / sigma^2) )

    References:
        1. |citation1|_

        .. _citation1: https://link.aps.org/doi/10.1103/PhysRevA.83.012308

        .. |citation1| replace:: *Gambetta, J. M., Motzoi, F., Merkel, S. T. & Wilhelm, F. K.
           Analytic control methods for high-fidelity unitary operations
           in a weakly nonlinear oscillator. Phys. Rev. A 83, 012308 (2011).*

        2. |citation2|_

        .. _citation2: https://link.aps.org/doi/10.1103/PhysRevLett.103.110501

        .. |citation2| replace:: *F. Motzoi, J. M. Gambetta, P. Rebentrost, and F. K. Wilhelm
           Phys. Rev. Lett. 103, 110501 – Published 8 September 2009.*
    """

    def __init__(self,
                 duration: int,
                 amp: complex,
                 sigma: float,
                 beta: float,
                 name: Optional[str] = None):
        """Initialize the drag pulse.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            amp: The amplitude of the Drag envelope.
            sigma: A measure of how wide or narrow the Gaussian peak is; described mathematically
                   in the class docstring.
            beta: The correction amplitude.
            name: Display name for this pulse envelope.
        """
        self._amp = complex(amp)
        self._sigma = sigma
        self._beta = beta
        super().__init__(duration=duration, name=name)

    @property
    def amp(self) -> complex:
        """The Gaussian amplitude."""
        return self._amp

    @property
    def sigma(self) -> float:
        """The Gaussian standard deviation of the pulse width."""
        return self._sigma

    @property
    def beta(self) -> float:
        """The weighing factor for the Gaussian derivative component of the waveform."""
        return self._beta

    def get_waveform(self) -> Waveform:
        return drag(duration=self.duration, amp=self.amp, sigma=self.sigma,
                    beta=self.beta, zero_ends=False)

    def validate_parameters(self) -> None:
        if abs(self.amp) > 1.:
            raise PulseError("The amplitude norm must be <= 1, "
                             "found: {}".format(abs(self.amp)))
        if self.sigma <= 0:
            raise PulseError("Sigma must be greater than 0.")
        if isinstance(self.beta, complex):
            raise PulseError("Beta must be real.")
        # Check if beta is too large: the amplitude norm must be <=1 for all points
        if self.beta > self.sigma:
            # If beta <= sigma, then the maximum amplitude is at duration / 2, which is
            # already constrainted by self.amp <= 1

            # 1. Find the first maxima associated with the beta * d/dx gaussian term
            #    This eq is derived from solving for the roots of the norm of the drag function.
            #    There is a second maxima mirrored around the center of the pulse with the same
            #    norm as the first, so checking the value at the first x maxima is sufficient.
            argmax_x = (self.duration / 2
                        - (self.sigma / self.beta) * math.sqrt(self.beta ** 2 - self.sigma ** 2))
            if argmax_x < 0:
                # If the max point is out of range, either end of the pulse will do
                argmax_x = 0

            # 2. Find the value at that maximum
            max_val = continuous.drag(np.array(argmax_x), sigma=self.sigma,
                                      beta=self.beta, amp=self.amp, center=self.duration / 2)
            if abs(max_val) > 1.:
                raise PulseError("Beta is too large; pulse amplitude norm exceeds 1.")

    @property
    def parameters(self) -> Dict[str, Any]:
        return {"duration": self.duration, "amp": self.amp, "sigma": self.sigma,
                "beta": self.beta}

    def __repr__(self) -> str:
        return "{}(duration={}, amp={}, sigma={}, beta={}{})" \
               "".format(self.__class__.__name__, self.duration, self.amp, self.sigma, self.beta,
                         ", name='{}'".format(self.name) if self.name is not None else "")

class Constant(ParametricPulse):
    """
    A simple constant pulse, with an amplitude value and a duration:

    .. math::

        f(x) = amp    ,  0 <= x < duration
        f(x) = 0      ,  elsewhere
    """

    def __init__(self,
                 duration: int,
                 amp: complex,
                 name: Optional[str] = None):
        """
        Initialize the constant-valued pulse.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            amp: The amplitude of the constant square pulse.
            name: Display name for this pulse envelope.
        """
        self._amp = complex(amp)
        super().__init__(duration=duration, name=name)

    @property
    def amp(self) -> complex:
        """The constant value amplitude."""
        return self._amp

    def get_waveform(self) -> Waveform:
        return constant(duration=self.duration, amp=self.amp)

    def validate_parameters(self) -> None:
        if abs(self.amp) > 1.:
            raise PulseError("The amplitude norm must be <= 1, "
                             "found: {}".format(abs(self.amp)))

    @property
    def parameters(self) -> Dict[str, Any]:
        return {"duration": self.duration, "amp": self.amp}

    def __repr__(self) -> str:
        return "{}(duration={}, amp={}{})" \
               "".format(self.__class__.__name__, self.duration, self.amp,
                         ", name='{}'".format(self.name) if self.name is not None else "")

class ConstantPulse(Constant):
    """
    Deprecated. A simple constant pulse, with an amplitude value and a duration:

    .. math::

        f(x) = amp    ,  0 <= x < duration
        f(x) = 0      ,  elsewhere
    """

    def __init__(self,
                 duration: int,
                 amp: complex,
                 name: Optional[str] = None):
        """
        Initialize the constant-valued pulse.

        Args:
            duration: Pulse length in terms of the the sampling period `dt`.
            amp: The amplitude of the constant square pulse.
            name: Display name for this pulse envelope.
        """
        super(ConstantPulse, self).__init__(duration, amp, name)
        warnings.warn("The ConstantPulse is deprecated. Use Constant instead", DeprecationWarning)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
THOR detects differential peaks in multiple ChIP-seq profiles associated
with two distinct biological conditions.

Copyright (C) 2014-2016 Manuel Allhoff (allhoff@aices.rwth-aachen.de)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

@author: Manuel Allhoff
"""

from __future__ import print_function
import string
from hmmlearn.hmm import _BaseHMM

import sys
from math import fabs
from scipy.special import logsumexp

import numpy as np
from .neg_bin import NegBin

import warnings

def _get_pvalue_distr(mu, alpha, tracker):
    """Derive NB1 parameters for p-value calculation"""
    mu = mu[0,0]
    alpha = alpha[0,0] / 10000.
    tracker.write(text=str(mu), header="Neg. Bin. distribution for p-value estimates (mu)")
    tracker.write(text=str(alpha), header="Neg. Bin. distribution for p-value estimates (alpha)")

    nb = NegBin(mu, alpha)
    return {'distr_name': 'nb', 'distr': nb}

def get_init_parameters(s0, s1, s2, **info):
    """For given training set (s0: Background, s1: Gaining, s2: loseing) get inital mu, alpha for NB1."""
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        mu = np.matrix([np.mean(map(lambda x: x[i], s)) for i in range(2) for s in [s0, s1, s2]]).reshape(2, 3)
        var = np.matrix([np.var(map(lambda x: x[i], s)) for i in range(2) for s in [s0, s1, s2]]).reshape(2, 3)

    alpha = (var - mu) / np.square(mu)

    # alpha[np.isnan(alpha)] = 0.001

    alpha[alpha < 0] = 0.001

    for el in [mu, alpha]:
        high = min(el[0,1], el[1,2]) + 0.5 * fabs(el[0,1] - el[1,2])
        low = min(el[1,1], el[0,2]) + 0.5 * fabs(el[1,1] - el[0,2])
        med = np.mean([el[0,0], el[1,0]])
        el[0,1] = high
        el[1,2] = high
        el[1,1] = low
        el[0,2] = low
        el[0,0] = med
        el[1,0] = med

    return alpha, mu

class NegBinRepHMM(_BaseHMM):
    def __init__(self, alpha, mu, dim_cond_1, dim_cond_2, init_state_seq=None, n_components=3, covariance_type='diag',
                  startprob_prior=1.0, transmat_prior=1.0, func=None,
                 algorithm="viterbi", means_prior=None, means_weight=0,
                 covars_prior=1e-2, covars_weight=1,
                 random_state=None, n_iter=30, thresh=1e-2,
                 params=string.ascii_letters,
                 init_params=string.ascii_letters):

        _BaseHMM.__init__(self, n_components,
                          startprob_prior=startprob_prior,
                          transmat_prior=transmat_prior, algorithm=algorithm,
                          random_state=random_state, n_iter=n_iter,
                          tol=thresh, params=params,
                          init_params=init_params)

        self.dim = [dim_cond_1, dim_cond_2] #dimension of one emission
        self.n_features = 2 #sum(self.dim) #emission dimension
        self.alpha = alpha
        self.mu = mu
        self._update_distr(self.mu, self.alpha)
        self.func = func
        self.em_prob = 0

    def fit(self, obs, three_para):
        """Estimate model parameters.

        An initialization step is performed before entering the EM
        algorithm. If you want to avoid this step, pass proper
        ``init_params`` keyword argument to estimator's constructor.

        Parameters
        ----------
        obs : list
            List of array-like observation sequences, each of which
            has shape (n_i, n_features), where n_i is the length of
            the i_th observation.

        Notes
        -----
        In general, `logprob` should be non-decreasing unless
        aggressive pruning is used.  Decreasing `logprob` is generally
        a sign of overfitting (e.g. a covariance parameter getting too
        small).  You can fix this by getting more training data,
        or strengthening the appropriate subclass-specific regularization
        parameter.
        """

        # what does this mean??
        self._init(obs, self.init_params)

        logprob = []
        for i in range(self.n_iter):
            # Expectation step
            stats = self._initialize_sufficient_statistics()
            curr_logprob = 0
            for seq in obs:
                framelogprob = self._compute_log_likelihood(seq)
                lpr, fwdlattice = self._do_forward_pass(framelogprob)
                bwdlattice = self._do_backward_pass(framelogprob)
                gamma = fwdlattice + bwdlattice
                posteriors = np.exp(gamma.T - logsumexp(gamma, axis=1)).T
                curr_logprob += lpr
                self._accumulate_sufficient_statistics(
                    stats, seq, framelogprob, posteriors, fwdlattice,
                    bwdlattice)
            logprob.append(curr_logprob)

            # Check for convergence.
            if i > 0 and logprob[-1] - logprob[-2] < self.tol:
                break

            # Maximization step
            self._do_mstep(stats, three_para)
        #print("Logprob of all M-steps: %s" %logprob, file=sys.stderr)
        self.em_prob = logprob[-1]
        return self

    def _update_distr(self, mu, alpha):
        """Update distributions assigned to each state with new mu and alpha"""
        raw1 = [NegBin(mu[0, 0], alpha[0, 0]), NegBin(mu[0, 1], alpha[0, 1]), NegBin(mu[0, 2], alpha[0, 2])]
        raw2 = [NegBin(mu[1, 0], alpha[1, 0]), NegBin(mu[1, 1], alpha[1, 1]), NegBin(mu[1, 2], alpha[1, 2])]

        self.neg_distr = np.matrix([raw1, raw2]) #matrix of all Neg. Bin. Distributions, columns=HMM's state (3), row=#samples (2)

    def get_alpha(self, m):
        """Return alpha for a given mu based on empirical variance"""
        var = self.func(m)
        try:
            return max((var - m) / m**2, 1e-300)
        except Warning:
            if m**2 > 1e-300:
                return max((var - m) / m**2, 1e-300)
            else:
                return 1e-300

    def _compute_log_likelihood(self, X):
        matrix = []
        lookup = {}
        for x in X: #over all observations
            row = []
            for i in range(self.n_components): #over number of HMM's state
                r_sum = 0
                for j in range(self.n_features): #over dim
                    it = range(self.dim[0]) if j == 0 else range(self.dim[0], self.dim[0] + self.dim[1]) #grab proper ob
                    for k in it:
                        index = (int(x[k]), i, j)
                        if lookup.has_key( index ):
                            r_sum += lookup[index]
                        else:
                            y = float(self.neg_distr[j,i].logpdf(x[k]))
                            lookup[index] = y
                            r_sum += y
                row.append(r_sum)

            matrix.append(row)
        return np.asarray(matrix)

    def _generate_sample_from_state(self, state, random_state=None):
        output = []
        for i, d in enumerate(self.dim):
            for _ in range(d):
                output.append( self.neg_distr[i,state].rvs() )

        return np.array(output)

    def _initialize_sufficient_statistics(self):
        stats = super(NegBinRepHMM, self)._initialize_sufficient_statistics()
        stats['post'] = np.zeros([self.n_features, self.n_components])
        stats['post_emission'] = np.zeros([self.n_features, self.n_components]) #dim X states

        return stats

    def _help_accumulate_sufficient_statistics(self, obs, stats, posteriors):
        for t, symbol in enumerate(obs):
            stats['post'][0] += posteriors[t]
            stats['post'][1] += posteriors[t]

            pot_it = [range(self.dim[0]), range(self.dim[0], self.dim[0] + self.dim[1])] #consider both classes
            for j, it in enumerate(pot_it):
                for i in it:
                    stats['post_emission'][j] += posteriors[t] * symbol[i]

        stats['post'][0] = stats['post'][0] * self.dim[0]
        stats['post'][1] = stats['post'][1] * self.dim[1]

        stats['posterior'] = np.copy(posteriors)

    def _valid_posteriors(self, posteriors, obs):

        warnings.filterwarnings('error')

        for i in range(len(obs)):
            state_1 = False
            c1, c2 = np.mean(obs[i][:self.dim[0]]), np.mean(obs[i][self.dim[0]:]) #counts of samples

            if posteriors[i][0] > 0.5:
                state_1 = True

            if c1 > c2: #state 1
                if fabs(posteriors[i][2] - 1) < 1e-200:
                    posteriors[i] = np.array([1, 0, 0])
                else:
                    if not state_1 and posteriors[i][2] > posteriors[i][1]:
                        try:
                            post_s2 = 0
                            post_s0 = posteriors[i][0] / (posteriors[i][0] + posteriors[i][1])
                            post_s1 = posteriors[i][1] / (posteriors[i][0] + posteriors[i][1])
                            posteriors[i] = np.array([post_s0, post_s1, post_s2])
                        except RuntimeWarning:
                            print(posteriors[i], c1, c2, file=sys.stderr)

            if c2 > c1: #state 2
                if fabs(posteriors[i][1] - 1) < 1e-200:
                    posteriors[i] = np.array([1, 0, 0])
                else:
                    if not state_1 and posteriors[i][1] > posteriors[i][2]:
                        try:
                            post_s1 = 0
                            post_s0 = posteriors[i][0] / (posteriors[i][0] + posteriors[i][2])
                            post_s2 = posteriors[i][2] / (posteriors[i][0] + posteriors[i][2])
                            posteriors[i] = np.array([post_s0, post_s1, post_s2])
                        except RuntimeWarning:
                            print(posteriors[i], c1, c2, file=sys.stderr)

        warnings.resetwarnings()

        return posteriors

    def _accumulate_sufficient_statistics(self, stats, obs, framelogprob,
                                      posteriors, fwdlattice, bwdlattice
                                      ):
        super(NegBinRepHMM, self)._accumulate_sufficient_statistics(
            stats, obs, framelogprob, posteriors, fwdlattice, bwdlattice,
            )
        posteriors = self._valid_posteriors(posteriors, obs)
        self._help_accumulate_sufficient_statistics(obs, stats, posteriors)

    def _do_mstep(self, stats, three_para):
        super(NegBinRepHMM, self)._do_mstep(stats)

        if three_para:
            self.mu[0,1] = (stats['post_emission'][0][1] + stats['post_emission'][1][2]) / (stats['post'][0][1] + stats['post'][1][2])
            self.mu[1,1] = (stats['post_emission'][1][1] + stats['post_emission'][0][2]) / (stats['post'][1][1] + stats['post'][0][2])
            self.mu[0,0] = (stats['post_emission'][0][0] + stats['post_emission'][1][0]) / (stats['post'][0][0] + stats['post'][1][0])

            self.mu[1,2] = self.mu[0,1]
            self.mu[0,2] = self.mu[1,1]
            self.mu[1,0] = self.mu[0,0]
        else:
            self.mu[0,1] = (stats['post_emission'][0][1] + stats['post_emission'][1][2]) / (stats['post'][0][1] + stats['post'][1][2])
            self.mu[1,1] = (stats['post_emission'][1][1] + stats['post_emission'][0][2] + stats['post_emission'][0][0] + stats['post_emission'][0][1]) / (stats['post'][1][1] + stats['post'][0][2] + stats['post'][0][0] + stats['post'][0][1])

            self.mu[0,0] = self.mu[1,1]
            self.mu[1,2] = self.mu[0,1]
            self.mu[0,2] = self.mu[1,1]
            self.mu[1,0] = self.mu[0,0]

        self.alpha = np.matrix([map(lambda m: self.get_alpha(m), np.asarray(self.mu[i])[0]) for i in range(self.n_features)])
        self._update_distr(self.mu, self.alpha)

    def merge_distr(self):
        f = self.count_s2 / float(self.count_s1 + self.count_s2) #TODO exp_data.

        for el in [self.mu, self.alpha]:
            high = min(el[0,1], el[1,2]) + f * fabs(el[0,1] - el[1,2])
            low = min(el[1,1], el[0,2]) + f * fabs(el[1,1] - el[0,2])
            med = np.mean([el[0,0], el[1,0]])
            el[0,1] = high
            el[1,2] = high
            el[1,1] = low
            el[0,2] = low
            el[0,0] = low #min(med, low)
            el[1,0] = low #min(med, low)

        self._update_distr(self.mu, self.alpha)

if __name__ == '__main__':
    alpha = np.matrix([[0.2, 0.2, 0.2], [0.2, 0.2, 0.2]])
    mu = np.matrix([[15.,100.,10.], [10.,10.,100.]])
    f = lambda x: 0.4*x**2 + 1

    dim_cond_1 = 2
    dim_cond_2 = 3
    mean = [0, 0]
    cov = [[1, 0], [0, 100]]
    X= np.random.multivariate_normal(mean, cov, size=5)

    m2 = NegBinRepHMM(alpha = alpha, mu = np.matrix([[50.,130.,110.], [60.,100.,120.]]), dim_cond_1 = dim_cond_1, dim_cond_2 = dim_cond_2, func=f)
    m2.fit([X], three_para=False)

    posteriors = m2.predict_proba(X)
    e = m2.predict(X)
    for i, el in enumerate(X):
        print(el, e[i], sep='\t', file=sys.stderr)
    #print(np.max(posteriors, axis=1))
    print(m2.mu)