#!/usr/bin/python

# Comments and/or additions are welcome (send e-mail to:
# robert.paton@chem.ox.ac.uk

#####################################
#        vib_scale_factors.py       #
#####################################
###  Written by:  Rob Paton #########
###  Last modified:  Mar 27, 2017 ###
#####################################

import numpy as np

import sys
if sys.version_info > (2, ):
    # Py3
    Str_char = "U%d"
else:
    # Py2
    Str_char = "S%d"

'''
Frequency scaling factors, taken from the beta 2 version of version 3 of The Trhular group database (https://t1.chem.umn.edu/freqscale/index.html
I. M. Alecu, J. Zheng, Y. Zhao, and D. G. Truhlar, J. Chem. Theory Comput. 6, 2872-2887 (2010).

The array is ordered as:
[level/basis set, zpe_fac, zpe_ref, zpe_meth, harm_fac, harm_ref, harm_meth, fund_fac, fund_ref, fund_meth]
where zpe_fac, harm_fac and fund_fac are the scaling factors for ZPEs, harmonic frequencies, and fundamentals, respectively.

The ref and meth elements refer to original references and method of determinantion by the Truhlar group. All information taken from
https://comp.chem.umn.edu/freqscale/version3b2.htm

Methods
D: The scale factor was directly obtained from the ZPVE15/10 or F38/10 databases given in Ref. 1.
C: The scale factor was obtained by applying a small systematic correction of -0.0025 to preexisting scale factor. The references for the preexisting (uncorrected) scale factors are given in Supporting Information of Ref. 1 and in Version 1 of this database
R: The scale factor was obtained via the Reduced Scale Factor Optimization Model described in Ref. 1. Briefly, this entails using the ZPE6 database for determining ZPE scale factors, and/or using the universal scale factor ratios of aF/ZPE = 0.974 and aH/ZPE = 1.014 to obtain the respective values for the scale factors for fundamental and harmonic frequencies.
'''
scaling_refs = np.array([("none"), ("I. M. Alecu, J. Zheng, Y. Zhao, and D. G. Truhlar, J. Chem. Theory Comput. 6, 2872-2887 (2010)."),("Y. Zhao and D. G. Truhlar, unpublished (2003), modified by systematic correction of -0.0025 by I. M. Alecu (2010)."), ("I. M. Alecu, unpublished (2011)."), ("J. Zheng, R. J. Rocha, M. Pelegrini, L. F. A. Ferrao, E. F. V. Carvalho, O. Roberto-Neto, F. B. C. Machado, and D. G. Truhlar, J. Chem. Phys. 136, 184310/1-10 (2012)."), ("J. Zheng and D. G. Truhlar, unpublished (2014)."), ("J. Bao and D. G. Truhlar, unpublished (2014)."), ("H. Yu, J. Zheng, and D. G. Truhlar, unpublished (2015)")])

scaling_data = np.array([("AM1" ,  0.948 ,  1 ,  "R" ,  0.961 ,  1 ,  "R" ,  0.923 ,  1 ,  "R" ),("B1B95/6-31+G(d,p)" ,  0.971 ,  1 ,  "C" ,  0.985 ,  1 ,  "R" ,  0.946 ,  1 ,  "R" ),("B1B95/MG3S" ,  0.973 ,  1 ,  "C" ,  0.987 ,  1 ,  "R" ,  0.948 ,  1 ,  "R" ),("B1LYP/MG3S" ,  0.978 ,  1 ,  "D" ,  0.994 ,  1 ,  "D" ,  0.955 ,  1 ,  "D" ),("B3LYP/6-31G(2df,2p)" ,  0.981 ,  1 ,  "C" ,  0.995 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("B3LYP/6-31G(d)" ,  0.977 ,  1 ,  "R" ,  0.991 ,  1 ,  "R" ,  0.952 ,  1 ,  "R" ),("B3LYP/aug-cc-pVTZ" ,  0.985 ,  3 ,  "R" ,  0.999 ,  3 ,  "R" ,  0.959 ,  3 ,  "R" ),("B3LYP/def2-TZVP" ,  0.985 ,  3 ,  "R" ,  0.999 ,  3 ,  "R" ,  0.959 ,  3 ,  "R" ),("B3LYP/ma-TZVP" ,  0.986 ,  1 ,  "R" ,  1.000 ,  1 ,  "R" ,  0.960 ,  1 ,  "R" ),("B3LYP/MG3S" ,  0.983 ,  1 ,  "D" ,  0.998 ,  1 ,  "D" ,  0.960 ,  1 ,  "D" ),("B3P86/6-31G(d)" ,  0.971 ,  1 ,  "R" ,  0.985 ,  1 ,  "R" ,  0.946 ,  1 ,  "R" ),("B3PW91/6-31G(d)" ,  0.972 ,  1 ,  "R" ,  0.986 ,  1 ,  "R" ,  0.947 ,  1 ,  "R" ),("B97-3/ma-TZVP" ,  0.975 ,  1 ,  "R" ,  0.989 ,  1 ,  "R" ,  0.950 ,  1 ,  "R" ),("B97-3/MG3S" ,  0.972 ,  1 ,  "D" ,  0.986 ,  1 ,  "D" ,  0.947 ,  1 ,  "D" ),("B98/def2-TZVP" ,  0.984 ,  1 ,  "R" ,  0.998 ,  1 ,  "R" ,  0.958 ,  1 ,  "R" ),("B98/ma-TZVP" ,  0.985 ,  1 ,  "R" ,  0.999 ,  1 ,  "R" ,  0.959 ,  1 ,  "R" ),("B98/MG3S" ,  0.982 ,  1 ,  "D" ,  0.995 ,  1 ,  "D" ,  0.956 ,  1 ,  "D" ),("BB1K/6-31+G(d,p)" ,  0.954 ,  1 ,  "C" ,  0.967 ,  1 ,  "R" ,  0.929 ,  1 ,  "R" ),("BB1K/MG3S" ,  0.957 ,  1 ,  "C" ,  0.970 ,  1 ,  "R" ,  0.932 ,  1 ,  "R" ),("BB95/6-31+G(d,p)" ,  1.011 ,  1 ,  "C" ,  1.025 ,  1 ,  "R" ,  0.985 ,  1 ,  "R" ),("BB95/MG3S" ,  1.012 ,  1 ,  "C" ,  1.026 ,  1 ,  "R" ,  0.986 ,  1 ,  "R" ),("BLYP/6-311G(df,p)" ,  1.013 ,  1 ,  "R" ,  1.027 ,  1 ,  "R" ,  0.987 ,  1 ,  "R" ),("BLYP/6-31G(d)" ,  1.009 ,  1 ,  "R" ,  1.023 ,  1 ,  "R" ,  0.983 ,  1 ,  "R" ),("BLYP/MG3S" ,  1.013 ,  1 ,  "D" ,  1.031 ,  1 ,  "D" ,  0.991 ,  1 ,  "D" ),("BMC-CCSD" ,  0.985 ,  1 ,  "D" ,  1.001 ,  1 ,  "D" ,  0.962 ,  1 ,  "D" ),("BMK/ma-TZVP" ,  0.972 ,  1 ,  "R" ,  0.986 ,  1 ,  "R" ,  0.947 ,  1 ,  "R" ),("BMK/MG3S" ,  0.971 ,  1 ,  "D" ,  0.984 ,  1 ,  "D" ,  0.945 ,  1 ,  "D" ),("BP86/6-31G(d)" ,  1.007 ,  1 ,  "R" ,  1.021 ,  1 ,  "R" ,  0.981 ,  1 ,  "R" ),("BP86/ma-TZVP" ,  1.014 ,  1 ,  "R" ,  1.028 ,  1 ,  "R" ,  0.988 ,  1 ,  "R" ),("BPW60/6-311+G(d,p)" ,  0.934 ,  2 ,  "C" ,  0.910 ,  2 ,  "R" ,  0.947 ,  2 ,  "R" ),("BPW63/MG3S" ,  0.923 ,  2 ,  "C" ,  0.899 ,  2 ,  "R" ,  0.936 ,  2 ,  "R" ),("CAM-B3LYP/ma-TZVP" ,  0.976 ,  1 ,  "R" ,  0.990 ,  1 ,  "R" ,  0.951 ,  1 ,  "R" ),("CCSD(T)/jul-cc-pVTZ" , 0.984 ,  1 ,  "R" ,  0.998 ,  1 ,  "R" ,  0.958 ,  1 ,  "R" ),("CCSD(T)/aug-cc-pVTZ" , 0.987 ,  1 ,  "R" ,  1.001 ,  1 ,  "R" ,  0.961 ,  1 ,  "R" ),("CCSD(T)-F12/jul-cc-pVTZ" , 0.981 ,  1 ,  "R" ,  0.995 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("CCSD(T)-F12a/cc-pVTZ-F12" , 0.984 ,  1 ,  "R" ,  0.998 ,  1 ,  "R" ,  0.958 ,  1 ,  "R" ),("CCSD/jul-cc-pVTZ" ,  0.973 ,  1 ,  "R" ,  0.987 ,  1 ,  "R" ,  0.948 ,  1 ,  "R" ),("CCSD-F12/jul-cc-pVTZ" ,  0.971 ,  1 ,  "R" ,  0.985 ,  1 ,  "R" ,  0.946 ,  1 ,  "R" ),("G96LYP80/6-311+G(d,p)" ,  0.911 ,  2 ,  "C" ,  0.887 ,  2 ,  "R" ,  0.924 ,  2 ,  "R" ),("G96LYP82/MG3S" ,  0.907 ,  2 ,  "C" ,  0.883 ,  2 ,  "R" ,  0.920 ,  2 ,  "R" ),("GAM/def2-TZVP" ,  0.980 ,  7 ,  "D" ,  0.994 ,  7 ,  "D" ,  0.955 ,  7 ,  "D" ),("GAM/ma-TZVP" ,  0.981 ,  7 ,  "D" ,  0.995 ,  7 ,  "D" ,  0.956 ,  7 ,  "D" ),("HF/3-21G" ,  0.919 ,  1 ,  "R" ,  0.932 ,  1 ,  "R" ,  0.895 ,  1 ,  "R" ),("HF/6-31+G(d)" ,  0.911 ,  1 ,  "R" ,  0.924 ,  1 ,  "R" ,  0.887 ,  1 ,  "R" ),("HF/6-31+G(d,p)" ,  0.915 ,  1 ,  "C" ,  0.928 ,  1 ,  "R" ,  0.891 ,  1 ,  "R" ),("HF/6-311G(d,p)" ,  0.920 ,  1 ,  "R" ,  0.933 ,  1 ,  "R" ,  0.896 ,  1 ,  "R" ),("HF/6-311G(df,p)" ,  0.920 ,  1 ,  "R" ,  0.933 ,  1 ,  "R" ,  0.896 ,  1 ,  "R" ),("HF/6-31G(d)" ,  0.909 ,  1 ,  "R" ,  0.922 ,  1 ,  "R" ,  0.885 ,  1 ,  "R" ),("HF/6-31G(d,p)" ,  0.913 ,  1 ,  "R" ,  0.926 ,  1 ,  "R" ,  0.889 ,  1 ,  "R" ),("HF/MG3S" ,  0.919 ,  1 ,  "D" ,  0.932 ,  1 ,  "D" ,  0.895 ,  1 ,  "D" ),("HFLYP/MG3S" ,  0.899 ,  1 ,  "D" ,  0.912 ,  1 ,  "D" ,  0.876 ,  1 ,  "D" ),("HSEh1PBE/ma-TZVP" ,  0.979 ,  1 ,  "R" ,  0.993 ,  1 ,  "R" ,  0.954 ,  1 ,  "R" ),("M05/aug-cc-pVTZ" ,  0.978 ,  1 ,  "R" ,  0.992 ,  1 ,  "R" ,  0.953 ,  1 ,  "R" ),("M05/def2-TZVP" ,  0.978 ,  3 ,  "R" ,  0.991 ,  3 ,  "R" ,  0.952 ,  3 ,  "R" ),("M05/ma-TZVP" ,  0.979 ,  1 ,  "R" ,  0.993 ,  1 ,  "R" ,  0.954 ,  1 ,  "R" ),("M05/maug-cc-pVTZ" ,  0.978 ,  1 ,  "R" ,  0.992 ,  1 ,  "R" ,  0.953 ,  1 ,  "R" ),("M05/MG3S" ,  0.977 ,  1 ,  "D" ,  0.989 ,  1 ,  "D" ,  0.951 ,  1 ,  "D" ),("M05-2X/6-31+G(d,p)" ,  0.961 ,  1 ,  "D" ,  0.974 ,  1 ,  "D" ,  0.936 ,  1 ,  "D" ),("M05-2X/aug-cc-pVTZ" ,  0.964 ,  1 ,  "R" ,  0.977 ,  1 ,  "R" ,  0.939 ,  1 ,  "R" ),("M05-2X/def2-TZVPP" ,  0.962 ,  1 ,  "D" ,  0.976 ,  1 ,  "D" ,  0.938 ,  1 ,  "D" ),("M05-2X/ma-TZVP" ,  0.965 ,  1 ,  "R" ,  0.979 ,  1 ,  "R" ,  0.940 ,  1 ,  "R" ),("M05-2X/maug-cc-pVTZ" ,  0.964 ,  1 ,  "R" ,  0.977 ,  1 ,  "R" ,  0.939 ,  1 ,  "R" ),("M05-2X/MG3S" ,  0.962 ,  1 ,  "D" ,  0.975 ,  1 ,  "D" ,  0.937 ,  1 ,  "D" ),("M06/6-31+G(d,p)" ,  0.980 ,  1 ,  "D" ,  0.989 ,  1 ,  "D" ,  0.950 ,  1 ,  "D" ),("M06/aug-cc-pVTZ" ,  0.984 ,  1 ,  "R" ,  0.998 ,  1 ,  "R" ,  0.958 ,  1 ,  "R" ),("M06/def2-TZVP" ,  0.982 ,  3 ,  "R" ,  0.996 ,  3 ,  "R" ,  0.956 ,  3 ,  "R" ),("M06/def2-TZVPP" ,  0.979 ,  1 ,  "D" ,  0.992 ,  1 ,  "D" ,  0.953 ,  1 ,  "D" ),("M06/ma-TZVP" ,  0.982 ,  1 ,  "R" ,  0.996 ,  1 ,  "R" ,  0.956 ,  1 ,  "R" ),("M06/maug-cc-pVTZ" ,  0.982 ,  1 ,  "R" ,  0.996 ,  1 ,  "R" ,  0.956 ,  1 ,  "R" ),("M06/MG3S" ,  0.981 ,  1 ,  "D" ,  0.994 ,  1 ,  "D" ,  0.955 ,  1 ,  "D" ),("M06-2X/6-31+G(d,p)" ,  0.967 ,  1 ,  "D" ,  0.979 ,  1 ,  "D" ,  0.940 ,  1 ,  "D" ),("M06-2X/6-311+G(d,p)" ,  0.970 ,  5 ,  "D" ,  0.983 ,  5 ,  "R" ,  0.944 ,  5 ,  "R" ),("M06-2X/6-311++G(d,p)" ,  0.970 ,  5 ,  "D" ,  0.983 ,  5 ,  "R" ,  0.944 ,  5 ,  "R" ),("M06-2X/aug-cc-pVTZ" ,  0.971 ,  1 ,  "D" ,  0.985 ,  1 ,  "D" ,  0.946 ,  1 ,  "D" ),("M06-2X/def2-TZVP" ,  0.971 ,  7 ,  "D" ,  0.984 ,  7 ,  "D" ,  0.946 ,  7 ,  "D" ),("M06-2X/def2-QZVP" ,  0.970 ,  7 ,  "D" ,  0.983 ,  7 ,  "D" ,  0.945 ,  7 ,  "D" ),("M06-2X/def2-TZVPP" ,  0.970 ,  1 ,  "D" ,  0.983 ,  1 ,  "D" ,  0.945 ,  1 ,  "D" ),("M06-2X/ma-TZVP" ,  0.972 ,  1 ,  "R" ,  0.986 ,  1 ,  "R" ,  0.947 ,  1 ,  "R" ),("M06/maug-cc-pVTZ" , 0.971 ,  1 ,  "D" ,  0.984 ,  1 ,  "D" ,  0.945 ,  1 ,  "D" ),("M06-2X/MG3S" ,  0.970 ,  1 ,  "D" ,  0.982 ,  1 ,  "D" ,  0.944 ,  1 ,  "D" ),("M06-HF/6-31+G(d,p)" ,  0.954 ,  1 ,  "D" ,  0.969 ,  1 ,  "D" ,  0.931 ,  1 ,  "D" ),("M06-HF/aug-cc-pVTZ" ,  0.961 ,  1 ,  "R" ,  0.974 ,  1 ,  "R" ,  0.936 ,  1 ,  "R" ),("M06-HF/def2-TZVPP" ,  0.958 ,  1 ,  "D" ,  0.970 ,  1 ,  "D" ,  0.932 ,  1 ,  "D" ),("M06-HF/ma-TZVP" ,  0.957 ,  1 ,  "R" ,  0.970 ,  1 ,  "R" ,  0.932 ,  1 ,  "R" ),("M06-HF/maug-cc-pVTZ" ,  0.959 ,  1 ,  "R" ,  0.972 ,  1 ,  "R" ,  0.934 ,  1 ,  "R" ),("M06-HF/MG3S" ,  0.955 ,  1 ,  "D" ,  0.967 ,  1 ,  "D" ,  0.930 ,  1 ,  "D" ),("M06-L/6-31+G(d,p)" ,  0.978 ,  1 ,  "D" ,  0.992 ,  1 ,  "D" ,  0.953 ,  1 ,  "D" ),("M06-L/aug-cc-pVTZ" ,  0.980 ,  1 ,  "R" ,  0.994 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("M06-L(DKH2)/aug-cc-pwcVTZ-DK" ,  0.985 ,  1 ,  "D" ,  0.999 ,  1 ,  "R" ,  0.959 ,  1 ,  "R" ),("M06-L/def2-TZVP" ,  0.976 ,  3 ,  "R" ,  0.990 ,  3 ,  "R" ,  0.951 ,  3 ,  "R" ),("M06-L/def2-TZVPP" ,  0.976 ,  1 ,  "D" ,  0.995 ,  1 ,  "D" ,  0.956 ,  1 ,  "D" ),("M06-L/ma-TZVP" ,  0.977 ,  1 ,  "R" ,  0.991 ,  1 ,  "R" ,  0.952 ,  1 ,  "R" ),("M06-L/maug-cc-pVTZ" ,  0.977 ,  1 ,  "R" ,  0.991 ,  1 ,  "R" ,  0.952 ,  1 ,  "R" ),("M06-L/MG3S" ,  0.978 ,  1 ,  "D" ,  0.996 ,  1 ,  "D" ,  0.958 ,  1 ,  "D" ),("M08-HX/6-31+G(d,p)" ,  0.972 ,  1 ,  "D" ,  0.983 ,  1 ,  "D" ,  0.944 ,  1 ,  "D" ),("M08-HX/aug-cc-pVTZ" ,  0.975 ,  1 ,  "R" ,  0.989 ,  1 ,  "R" ,  0.950 ,  1 ,  "R" ),("M08-HX/cc-pVTZ+" ,  0.974 ,  1 ,  "D" ,  0.985 ,  1 ,  "D" ,  0.946 ,  1 ,  "D" ),("M08-HX/def2-TZVPP" ,  0.973 ,  1 ,  "D" ,  0.984 ,  1 ,  "D" ,  0.945 ,  1 ,  "D" ),("M08-HX/jun-cc-pVTZ" ,  0.974 ,  6 ,  "D" ,  0.986 ,  6 ,  "D" ,  0.947 ,  6 ,  "D" ),("M08-HX/ma-TZVP" ,  0.976 ,  1 ,  "R" ,  0.990 ,  1 ,  "R" ,  0.951 ,  1 ,  "R" ),("M08-HX/maug-cc-pVTZ" ,  0.976 ,  1 ,  "R" ,  0.990 ,  1 ,  "R" ,  0.951 ,  1 ,  "R" ),("M08-HX/MG3S" ,  0.973 ,  1 ,  "D" ,  0.984 ,  1 ,  "D" ,  0.946 ,  1 ,  "D" ),("M08-SO/6-31+G(d,p)" ,  0.979 ,  1 ,  "D" ,  0.989 ,  1 ,  "D" ,  0.951 ,  1 ,  "D" ),("M08-SO/aug-cc-pVTZ" ,  0.985 ,  1 ,  "R" ,  0.999 ,  1 ,  "R" ,  0.959 ,  1 ,  "R" ),("M08-SO/cc-pVTZ+" ,  0.982 ,  1 ,  "D" ,  0.995 ,  1 ,  "D" ,  0.956 ,  1 ,  "D" ),("M08-SO/def2-TZVPP" ,  0.980 ,  1 ,  "D" ,  0.993 ,  1 ,  "D" ,  0.954 ,  1 ,  "D" ),("M08-SO/ma-TZVP" ,  0.984 ,  1 ,  "R" ,  0.998 ,  1 ,  "R" ,  0.958 ,  1 ,  "R" ),("M08-SO/maug-cc-pVTZ" ,  0.983 ,  1 ,  "R" ,  0.997 ,  1 ,  "R" ,  0.957 ,  1 ,  "R" ),("M08-SO/MG3" ,  0.984 ,  4 ,  "D" ,  0.998 ,  4 ,  "R" ,  0.959 ,  4 ,  "R" ),("M08-SO/MG3S" ,  0.983 ,  1 ,  "D" ,  0.995 ,  1 ,  "D" ,  0.956 ,  1 ,  "D" ),("M08-SO/MG3SXP" ,  0.984 ,  1 ,  "D" ,  0.996 ,  1 ,  "D" ,  0.957 ,  1 ,  "D" ),("MN12-L/MG3S" ,  0.968 ,  6 ,  "D" ,  0.981 ,  6 ,  "D" ,  0.943 ,  6 ,  "D" ),("MN12-SX/6-311++G(d,p)" ,  0.976 ,  6 ,  "D" ,  0.986 ,  6 ,  "D" ,  0.947 ,  6 ,  "D" ),("MN15-L/MG3S" ,  0.977 ,  1 ,  "D" ,  0.991 ,  1 ,  "R" ,  0.952 ,  1 ,  "R" ),("MN15-L/maug-cc-pVTZ" ,  0.979 ,  1 ,  "D" ,  0.993 ,  1 ,  "R" ,  0.954 ,  1 ,  "R" ),("MC3BB" ,  0.965 ,  1 ,  "C" ,  0.979 ,  1 ,  "R" ,  0.940 ,  1 ,  "R" ),("MC3MPW" ,  0.964 ,  1 ,  "C" ,  0.977 ,  1 ,  "R" ,  0.939 ,  1 ,  "R" ),("MC-QCISD/3" ,  0.992 ,  1 ,  "C" ,  1.006 ,  1 ,  "R" ,  0.966 ,  1 ,  "R" ),("MOHLYP/ma-TZVP" ,  1.027 ,  1 ,  "R" ,  1.041 ,  1 ,  "R" ,  1.000 ,  1 ,  "R" ),("MOHLYP/MG3S" ,  1.022 ,  1 ,  "R" ,  1.036 ,  1 ,  "R" ,  0.995 ,  1 ,  "R" ),("MP2(FC)/6-31+G(d,p)" , 0.968 ,  1 ,  "C" ,  0.982 ,  1 ,  "R" ,  0.943 ,  1 ,  "R" ),("MP2(FC)/6-311G(d,p)" , 0.970 ,  1 ,  "R" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("MP2(FC)/6-31G(d)" ,  0.964 ,  1 ,  "R" ,  0.977 ,  1 ,  "R" ,  0.939 ,  1 ,  "R" ),("MP2(FC)/6-31G(d,p)" , 0.958 ,  1 ,  "R" ,  0.971 ,  1 ,  "R" ,  0.933 ,  1 ,  "R" ),("MP2(FC)/cc-pVDZ" ,  0.977 ,  1 ,  "C" ,  0.991 ,  1 ,  "R" ,  0.952 ,  1 ,  "R" ),("MP2(FULL)/6-31G(d)" ,  0.963 ,  1 ,  "R" ,  0.976 ,  1 ,  "R" ,  0.938 ,  1 ,  "R" ),("MP4(SDQ)/jul-cc-pVTZ", 0.973 ,  1 ,  "R" ,  0.987 ,  1 ,  "R" ,  0.948 ,  1 ,  "R" ),("MPW1B95/6-31+G(d,p)" ,  0.970 ,  1 ,  "C" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("MPW1B95/MG3" ,  0.970 ,  1 ,  "C" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("MPW1B95/MG3S" ,  0.972 ,  1 ,  "C" ,  0.986 ,  1 ,  "R" ,  0.947 ,  1 ,  "R" ),("MPW1K/6-31+G(d,p)" ,  0.949 ,  1 ,  "C" ,  0.962 ,  1 ,  "R" ,  0.924 ,  1 ,  "R" ),("MPW1K/ma-TZVP" ,  0.956 ,  1 ,  "R" ,  0.969 ,  1 ,  "R" ,  0.931 ,  1 ,  "R" ),("MPW1K/MG3" ,  0.953 ,  1 ,  "C" ,  0.966 ,  1 ,  "R" ,  0.928 ,  1 ,  "R" ),("MPW1K/MG3S" ,  0.956 ,  1 ,  "C" ,  0.969 ,  1 ,  "R" ,  0.931 ,  1 ,  "R" ),("MPW1K/MIDI!" ,  0.953 ,  1 ,  "R" ,  0.966 ,  1 ,  "R" ,  0.928 ,  1 ,  "R" ),("MPW1K/MIDIY" ,  0.947 ,  1 ,  "R" ,  0.960 ,  1 ,  "R" ,  0.922 ,  1 ,  "R" ),("MPW3LYP/6-31+G(d,p)" ,  0.980 ,  1 ,  "C" ,  0.994 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("MPW3LYP/6-311+G(2d,p)" ,  0.986 ,  1 ,  "R" ,  1.000 ,  1 ,  "R" ,  0.960 ,  1 ,  "R" ),("MPW3LYP/6-31G(d)" ,  0.976 ,  1 ,  "R" ,  0.990 ,  1 ,  "R" ,  0.951 ,  1 ,  "R" ),("MPW3LYP/ma-TZVP" ,  0.986 ,  1 ,  "R" ,  1.000 ,  1 ,  "R" ,  0.960 ,  1 ,  "R" ),("MPW3LYP/MG3S" ,  0.982 ,  1 ,  "C" ,  0.996 ,  1 ,  "R" ,  0.956 ,  1 ,  "R" ),("MPW74/6-311+G(d,p)" ,  0.912 ,  2 ,  "C" ,  0.888 ,  2 ,  "R" ,  0.925 ,  2 ,  "R" ),("MPW76/MG3S" ,  0.909 ,  2 ,  "C" ,  0.885 ,  2 ,  "R" ,  0.922 ,  2 ,  "R" ),("MPWB1K/6-31+G(d,p)" ,  0.951 ,  1 ,  "C" ,  0.964 ,  1 ,  "R" ,  0.926 ,  1 ,  "R" ),("MPWB1K/MG3S" ,  0.954 ,  1 ,  "C" ,  0.967 ,  1 ,  "R" ,  0.929 ,  1 ,  "R" ),("MPWLYP1M/ma-TZVP" ,  1.009 ,  1 ,  "R" ,  1.023 ,  1 ,  "R" ,  0.983 ,  1 ,  "R" ),("OreLYP/ma-TZVP" ,  1.010 ,  7 ,  "D" ,  1.024 ,  7 ,  "D" ,  0.984 ,  7 ,  "D" ),("OreLYP/def2-TZVP" ,  1.008 ,  7 ,  "D" ,  1.023 ,  7 ,  "D" ,  0.982 ,  7 ,  "D" ),("PBE/def2-TZVP" ,  1.011 ,  3 ,  "R" ,  1.026 ,  3 ,  "R" ,  0.985 ,  3 ,  "R" ),("PBE/MG3S" ,  1.010 ,  1 ,  "D" ,  1.025 ,  1 ,  "D" ,  0.985 ,  1 ,  "D" ),("PBE/ma-TZVP" ,  1.014 ,  1 ,  "D" ,  1.028 ,  1 ,  "D" ,  0.987 ,  1 ,  "D" ),("PBE0/MG3S" ,  0.975 ,  1 ,  "D" ,  0.989 ,  1 ,  "D" ,  0.950 ,  1 ,  "D" ),("PBE1KCIS/MG3" ,  0.981 ,  1 ,  "C" ,  0.995 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("PBE1KCIS/MG3S" ,  0.981 ,  1 ,  "C" ,  0.995 ,  1 ,  "R" ,  0.955 ,  1 ,  "R" ),("PM3" ,  0.940 ,  1 ,  "R" ,  0.953 ,  1 ,  "R" ,  0.916 ,  1 ,  "R" ),("PM6" ,  1.078 ,  1 ,  "R" ,  1.093 ,  1 ,  "R" ,  1.050 ,  1 ,  "R" ),("PW6B95/6-31+G(d,p)" ,  0.970 ,  1 ,  "C" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("QCISD(FC)/6-31G(d)" ,  0.973 ,  1 ,  "R" ,  0.987 ,  1 ,  "R" ,  0.948 ,  1 ,  "R" ),("revTPSS/def2-TZVP" ,  0.998 ,  7 ,  "D" ,  1.012 ,  7 ,  "D" ,  0.972 ,  7 ,  "D" ),("revTPSS/ma-TZVP" ,  0.999 ,  7 ,  "D" ,  1.013 ,  7 ,  "D" ,  0.973 ,  7 ,  "D" ),("SOGGA/ma-TZVP" ,  1.017 ,  1 ,  "R" ,  1.031 ,  1 ,  "R" ,  0.991 ,  1 ,  "R" ),("tHCTHhyb/ma-TZVP" ,  0.989 ,  1 ,  "R" ,  1.003 ,  1 ,  "R" ,  0.963 ,  1 ,  "R" ),("TPSS1KCIS/def2-TZVP" ,  0.982 ,  1 ,  "R" ,  0.996 ,  1 ,  "R" ,  0.956 ,  1 ,  "R" ),("TPSS1KCIS/ma-TZVP" ,  0.983 ,  1 ,  "R" ,  0.997 ,  1 ,  "R" ,  0.957 ,  1 ,  "R" ),("TPSSh/MG3S" ,  0.984 ,  1 ,  "D" ,  1.002 ,  1 ,  "D" ,  0.963 ,  1 ,  "D" ),("VSXC/MG3S" ,  0.986 ,  1 ,  "D" ,  1.001 ,  1 ,  "D" ,  0.962 ,  1 ,  "D" ),("wB97/def2-TZVP", 0.969 ,  1 ,  "R" ,  0.983 ,  1 ,  "R" ,  0.944 ,  1 ,  "R" ),("wB97/ma-TZVP", 0.970 ,  1 ,  "R" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("wB97X/def2-TZVP", 0.970 ,  1 ,  "R" ,  0.984 ,  1 ,  "R" ,  0.945 ,  1 ,  "R" ),("wB97X/ma-TZVP" , 0.971 ,  1 ,  "R" ,  0.985 ,  1 ,  "R" ,  0.946 ,  1 ,  "R" ),("wB97X-D/def2-TZVP", 0.975 ,  1 ,  "R" ,  0.989 ,  1 ,  "R" ,  0.950 ,  1 ,  "R" ),("wB97X-D/ma-TZVP", 0.975 ,  1 ,  "R" ,  0.989 ,  1 ,  "R" ,  0.950 ,  1 ,  "R" ),("wB97X-D/maug-cc-pVTZ", 0.974 ,  1 ,  "R" ,  0.988 ,  1 ,  "R" ,  0.949 ,  1 ,  "R" ),("X1B95/6-31+G(d,p)" ,  0.968 ,  1 ,  "C" ,  0.982 ,  1 ,  "R" ,  0.943 ,  1 ,  "R" ),("X1B95/MG3S" ,  0.971 ,  1 ,  "C" ,  0.985 ,  1 ,  "R" ,  0.946 ,  1 ,  "R" ),("XB1K/6-31+G(d,p)" ,  0.952 ,  1 ,  "C" ,  0.965 ,  1 ,  "R" ,  0.927 ,  1 ,  "R" ),("XB1K/MG3S" ,  0.955 ,  1 ,  "C" ,  0.968 ,  1 ,  "R" ,  0.930 ,  1 ,  "R" )],dtype=[('level', Str_char%30),('zpe_fac', 'f4'), ('zpe_ref', 'i4'), ('zpe_meth', Str_char%1), ('harm_fac', 'f4'), ('harm_ref', 'i4'), ('harm_meth', Str_char%1), ('fund_fac', 'f4'), ('fund_ref', 'i4'), ('fund_meth', Str_char%1)])


import numpy as np
import torch
import json
import os
from ops.merger import OverlapCubeMerger
from ops.cube import ActivityTypeMEVA,CubeActivities,ActivityTypeVIRAT
import decord
from decord import VideoReader

def get_obj_types(event):
    toks = event.split("_")
    if "person" in toks and "vehicle" in toks :
        return ["Vehicle","Person"]
    elif "person" in toks:
        return ["Person"]
    elif "vehicle" in toks:
        return ["Vehicle"]
    else:
        print (event)
        raise RuntimeError("unknown events")

def get_file_index(filesProcessed):
    new_dict = {}
    for f in filesProcessed:
        new_dict[f]={"framerate": 30.0, "selected": {"0": 1, "9000": 0}}
    return new_dict

def get_activity_index(activities):
    new_dict = {}
    for act in activities:
        new_dict[act] = {"objectTypes":get_obj_types(act)}
    return new_dict

def getoutput(vid_names,video_pred_topall,video_prob_topall,event_dict,args,pred_frames = 64,filter_rate = 0.4):
    new_dict = {}
    for name, pred_all, prob_all in zip(vid_names, video_pred_topall,video_prob_topall):
        if args.dataset in ["MASK","DET","MEVA"]:
            video = name.split("_")[0]
        else:
            video = "_".join(name.split("_")[:-1])
        prop = name.split("_")[-1]
        if video not in list(new_dict.keys()):
            new_dict[video] = {}
            for event in event_dict:
                new_dict[video][event] = []
            new_dict[video]["props"] = []
        new_dict[video]["props"].append(prop)
        for i in range(len(pred_all)):
            pred = pred_all[i]
            event = event_dict[int(pred)]
            prob = prob_all[i]
            new_dict[video][event].append(str(prob))
    videos = list(new_dict.keys())
    database = json.load(open(args.val_list,"r"))["database"]
    props = os.listdir(os.path.join(args.prop_path,args.tst_name))
    merger = OverlapCubeMerger()
    activities = []
    files = []
    for prop in props:
        vid_name = prop.strip(".csv")
        if vid_name+args.vid_type not in files:
            files.append(vid_name+args.vid_type )
    for vid_name in videos:
        cube_tensor = torch.zeros(len(new_dict[vid_name]["props"])*len(event_dict),9)
        for i in range(len(new_dict[vid_name]["props"])):
            prop_name = vid_name+"_"+new_dict[vid_name]["props"][i]
            t0 = float(database[prop_name]["annotations"]["start"])
            t1 = float(database[prop_name]["annotations"]["end"])
            tid = int(database[prop_name]["tid"])
            x0,y0,x1,y1 = database[prop_name]["annotations"]["bbox"]
            for j in range(len(event_dict)):
                score = float(new_dict[vid_name][event_dict[j]][i])
                cube_tensor[i*len(event_dict)+j] = torch.tensor([tid,j+1,score,t0,t1,x0,y0,x1,y1])
        if args.dataset in ["MEVA","MASK","DET"]:
            cube_acts = CubeActivities(cube_tensor, vid_name+args.vid_type , ActivityTypeMEVA)
        else:
            cube_acts = CubeActivities(cube_tensor, vid_name+args.vid_type , ActivityTypeVIRAT)
        filtered_acts = merger(cube_acts)
        acts = filtered_acts.to_official()
        for act in acts:
            act["activityID"] = event_dict.index(act["activity"])
            activities.append(act)
    res_dict = {"filesProcessed":files,"activities":[]}
    file_dict = get_file_index(res_dict["filesProcessed"])
    eve_dict = get_activity_index(event_dict)
    sum_frame = 0
    fnames = file_dict.keys()
    for vid in fnames:
        vr = VideoReader(os.path.join(args.root_path,vid))
        sum_frame+=len(vr)
    loc = int(sum_frame/pred_frames*filter_rate)
    scores = [0]*len(event_dict)
    rank_lst = [[]]*len(event_dict)
    for act in activities:
        rank_lst[act["activityID"]].append(act["presenceConf"])
    for i in range(len(event_dict)):
        rank_lst[i].sort()
        if loc >=len(rank_lst[i]):
            loci = len(rank_lst[i])-1
        else:
            loci = loc
        scores[i] = rank_lst[i][-loci]

    for act in activities:
        if act["presenceConf"]>= scores[act["activityID"]]:
            res_dict["activities"].append(act)

    return res_dict,file_dict,eve_dict

def get_eventdict(args):
    if args.dataset in ["DET","MASK","MEVA"]:
        f = open("./labels_det.txt","r")
    elif args.dataset in ["VIRAT"]:
        f = open("./labels_virat.txt","r")
    events = f.readlines()
    event_dict = []
    for event in events:
        event_dict.append(event.strip())
    print(event_dict)
    return event_dict

def softmax(scores):
    es = np.exp(scores - scores.max(axis=-1)[..., None])
    return es / es.sum(axis=-1)[..., None]

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy_nll(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

def accuracy_bce(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    # print(output.size())
    # print(target.size())
    inner_target = []
    inner_output = []
    batch_size = target.size(0)
    for i in range(batch_size):
        if torch.max(target[i])>0:
            inner_output.append(output[i])
            inner_target.append(target[i])
    if len(inner_target)==0:
        return [torch.tensor(-1).cuda(),torch.tensor(-1).cuda()],0
    inner_output = torch.stack(inner_output)
    inner_target = torch.stack(inner_target)
    inner_target = torch.argmax(inner_target,dim=1)

    _, pred = inner_output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(inner_target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res, len(inner_target)


# %matplotlib inline
import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.utils.data import DataLoader,Dataset
import matplotlib.pyplot as plt
import torchvision.utils
import numpy as np
import random
from PIL import Image
import torch
from torch.autograd import Variable
import PIL.ImageOps
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from model import SiameseNetwork, ContrastiveLoss
import sys
from config import Config

def imshow(img,text=None,should_save=False):
    npimg = img.numpy()
    plt.axis("off")
    if text:
        plt.text(75, 8, text, style='italic',fontweight='bold',
            bbox={'facecolor':'white', 'alpha':0.8, 'pad':10})
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

def show_plot(iteration,loss):
    plt.plot(iteration,loss)
    plt.show()

class SiameseNetworkDataset(Dataset):

    def __init__(self, imageFolderDataset, transform=None, should_invert=True):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert

    def __get_imgs(self):
        # we need to make sure approx 50% of images are in the same class
        if random.randint(0, 1):
            the_class = random.choice(self.imageFolderDataset.classes)
            img0_tuple, img1_tuple = random.sample([x for x in self.imageFolderDataset.imgs if the_class in x[0]], 2)
        else:
            class_1, class_2 = random.sample(self.imageFolderDataset.classes, 2)
            img0_tuple = random.choice([x for x in self.imageFolderDataset.imgs if class_1 in x[0]])
            img1_tuple = random.choice([x for x in self.imageFolderDataset.imgs if class_2 in x[0]])
        return img0_tuple, img1_tuple

    def __getitem__(self, index):
        img0_tuple, img1_tuple = self.__get_imgs()

        img0 = Image.open(img0_tuple[0])
        img1 = Image.open(img1_tuple[0])
        img0 = img0.convert("L")
        img1 = img1.convert("L")

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)

        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)

        return img0, img1, torch.from_numpy(np.array([int(img1_tuple[1] != img0_tuple[1])], dtype=np.float32))

    def __len__(self):
        return len(self.imageFolderDataset.imgs)

def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press ⌘F8 to toggle the breakpoint.

def train():
    siamese_dataset = SiameseNetworkDataset(imageFolderDataset=folder_dataset,
                                            transform=transforms.Compose([transforms.Resize((100, 100)),
                                                                          transforms.ToTensor()
                                                                          ])
                                            , should_invert=False)

    train_dataloader = DataLoader(siamese_dataset,
                                  shuffle=True,
                                  num_workers=8,
                                  batch_size=Config.train_batch_size)

    global use_gpu
    if use_gpu:
        net = SiameseNetwork().cuda()
    else:
        net = SiameseNetwork().cpu()
    criterion = ContrastiveLoss()
    optimizer = optim.Adam(net.parameters(),lr = 0.0005 )

    counter = []
    loss_history = []
    iteration_number= 0

    for epoch in range(0, Config.max_epochs):
        for i, data in enumerate(train_dataloader,0):
            img0, img1 , label = data
            if use_gpu:
                img0, img1 , label = img0.cuda(), img1.cuda() , label.cuda()
            else:
                img0, img1 , label = img0.cpu(), img1.cpu() , label.cpu()
            optimizer.zero_grad()
            output1,output2 = net(img0,img1)
            loss_contrastive = criterion(output1,output2,label)
            loss_contrastive.backward()
            optimizer.step()
            if i %10 == 0 :
                print("Epoch number {}\n Current loss {}\n".format(epoch,loss_contrastive.item()))
                iteration_number +=10
                counter.append(iteration_number)
                loss_history.append(loss_contrastive.item())
    show_plot(counter,loss_history)

global use_gpu
use_gpu = False

if __name__ == '__main__':
    if len(sys.argv) > 1:
        use_gpu = "gpu" == sys.argv[1]
    print('use_gpu：', use_gpu)

    folder_dataset = dset.ImageFolder(root=Config.train_dir)
    siamese_dataset = SiameseNetworkDataset(imageFolderDataset=folder_dataset,
                                            transform=transforms.Compose([transforms.Resize((100, 100)),
                                                                          transforms.ToTensor()
                                                                          ])
                                            , should_invert=False)

    vis_dataloader = DataLoader(siamese_dataset,
                                shuffle=True,
                                num_workers=8,
                                batch_size=8)
    dataiter = iter(vis_dataloader)

    example_batch = next(dataiter)
    concatenated = torch.cat((example_batch[0], example_batch[1]), 0)
    imshow(torchvision.utils.make_grid(concatenated))
    print(example_batch[2].numpy())

# See PyCharm help at https://www.jetbrains.com/help/pycharm/


from itertools import product as product
from math import ceil

import matplotlib.pyplot as plt
import numpy as np
import torch

from utils.config import cfg_mnet

#-----------------------------#
#   中心解码，宽高解码
#-----------------------------#
def decode(loc, priors, variances):
    boxes = torch.cat((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
                    priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes

#-----------------------------#
#   关键点解码
#-----------------------------#
def decode_landm(pre, priors, variances):
    landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
                        ), dim=1)
    return landms

class Anchors(object):
    def __init__(self, cfg, image_size=None):
        super(Anchors, self).__init__()
        self.min_sizes  = cfg['min_sizes']
        self.steps      = cfg['steps']
        self.clip       = cfg['clip']
        #---------------------------#
        #   图片的尺寸
        #---------------------------#
        self.image_size = image_size
        #---------------------------#
        #   三个有效特征层高和宽
        #---------------------------#
        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]

    def get_anchors(self):
        anchors = []
        for k, f in enumerate(self.feature_maps):
            min_sizes = self.min_sizes[k]
            #-----------------------------------------#
            #   对特征层的高和宽进行循环迭代
            #-----------------------------------------#
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self.image_size[1]
                    s_ky = min_size / self.image_size[0]
                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        output = torch.Tensor(anchors).view(-1, 4)
        output  = np.zeros_like(anchors[:, :4])
        output[:,0] = anchors[:,0] - anchors[:,2] / 2
        output[:,1] = anchors[:,1] - anchors[:,3] / 2
        output[:,2] = anchors[:,0] + anchors[:,2] / 2
        output[:,3] = anchors[:,1] + anchors[:,3] / 2

        if self.clip:
            output = np.clip(output, 0, 1)
        return output

if __name__ == "__main__":
    cfg_mnet['image_size'] = 640
    #--------------------------------#
    #   先验框的生成
    #--------------------------------#
    cfg     = cfg_mnet
    anchors = Anchors(cfg, image_size = (cfg_mnet['image_size'], cfg_mnet['image_size'])).get_anchors()
    anchors = anchors[-800:] * cfg_mnet['image_size']

    #--------------------------------#
    #   先验框中心绘制
    #--------------------------------#
    center_x = (anchors[:, 0] + anchors[:, 2]) / 2
    center_y = (anchors[:, 1] + anchors[:, 3]) / 2

    fig = plt.figure()
    ax  = fig.add_subplot(121)
    plt.ylim(-300,900)
    plt.xlim(-300,900)
    ax.invert_yaxis()
    plt.scatter(center_x,center_y)

    #--------------------------------#
    #   先验框宽高绘制
    #--------------------------------#
    box_widths  = anchors[0:2,2] - anchors[0:2,0]
    box_heights = anchors[0:2,3] - anchors[0:2,1]
    for i in [0,1]:
        rect = plt.Rectangle([anchors[i, 0], anchors[i, 1]], box_widths[i], box_heights[i], color="r", fill=False)
        ax.add_patch(rect)

    #--------------------------------#
    #   先验框中心绘制
    #--------------------------------#
    ax = fig.add_subplot(122)
    plt.ylim(-300,900)
    plt.xlim(-300,900)
    ax.invert_yaxis()  #y轴反向
    plt.scatter(center_x,center_y)

    #--------------------------------#
    #   对先验框调整获得预测框
    #--------------------------------#
    mbox_loc = np.random.randn(800, 4)
    mbox_ldm = np.random.randn(800, 10)

    anchors[:, :2] = (anchors[:, :2] + anchors[:, 2:]) / 2
    anchors[:, 2:] = (anchors[:, 2:] - anchors[:, :2]) * 2

    mbox_loc                = torch.Tensor(mbox_loc)
    anchors                 = torch.Tensor(anchors)
    cfg_mnet['variance']    = torch.Tensor(cfg_mnet['variance'])
    decode_bbox             = decode(mbox_loc, anchors, cfg_mnet['variance'])

    box_widths  = decode_bbox[0: 2, 2] - decode_bbox[0: 2, 0]
    box_heights = decode_bbox[0: 2, 3] - decode_bbox[0: 2, 1]

    for i in [0,1]:
        rect = plt.Rectangle([decode_bbox[i, 0], decode_bbox[i, 1]], box_widths[i], box_heights[i], color="r", fill=False)
        plt.scatter((decode_bbox[i, 2] + decode_bbox[i, 0]) / 2, (decode_bbox[i,3] + decode_bbox[i,1]) / 2, color="b")
        ax.add_patch(rect)

    plt.show()


# test_model.py

import numpy as np
from grabscreen import grab_screen
import cv2
import time
from directkeys import PressKey,ReleaseKey, W, A, S, D
from alexnet import alexnet
from getkeys import key_check
from motion import motion_detection
from statistics import mode, mean
import random
from collections import deque, Counter

WIDTH = 160
HEIGHT = 120
LR = 1e-3
EPOCHS = 50
MODEL_NAME = 'pygta5-car-fast-0.001-alexnetv2-50-epochs-300K-data.model'.format(LR, 'alexnetv2',EPOCHS)

how_far_remove = 800
rs = (20, 15)
log_len = 25

motion_req = 800
motion_log = deque(maxlen = log_len)

t_time = 0.09

def straight():
    PressKey(W)
    ReleaseKey(A)
    ReleaseKey(D)
    ReleaseKey(S)

def left():
    if random.randrange(0, 3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    PressKey(A)
    ReleaseKey(S)
    ReleaseKey(D)
##    ReleaseKey(S)

def right():
    if random.randrange(0, 3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    PressKey(D)
    ReleaseKey(A)
    ReleaseKey(S)

def reverse():
    PressKey(S)
    ReleaseKey(A)
    ReleaseKey(W)
    ReleaseKey(D)

def forward_left():
    PressKey(W)
    PressKey(A)
    ReleaseKey(D)
    ReleaseKey(S)

def forward_right():
    PressKey(W)
    PressKey(D)
    ReleaseKey(A)
    ReleaseKey(S)

def reverse_left():
    PressKey(S)
    PressKey(A)
    ReleaseKey(W)
    ReleaseKey(D)

def reverse_right():
    PressKey(S)
    PressKey(D)
    ReleaseKey(W)
    ReleaseKey(A)

def no_keys():
    if random.randrange(0, 3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    ReleaseKey(A)
    ReleaseKey(S)
    ReleaseKey(D)

model = alexnet(WIDTH, HEIGHT, LR)
model.load(MODEL_NAME)

def main():
    last_time = time.time()
    for i in list(range(4))[::-1]:
        print(i+1)
        time.sleep(1)

    paused = False
    mode_choice = 0
    while(True):

        if not paused:
            # 800x600 windowed mode
            #screen =  np.array(ImageGrab.grab(bbox=(0,40,800,640)))
            screen = grab_screen(region=(0, 40, 1600, 910))
            print('loop took {} seconds'.format(time.time()-last_time))
            last_time = time.time()
            screen = cv2.cvtColor(screen, cv2.COLOR_BGR2RGB)
            screen = cv2.resize(screen, (160,120))
            prev = cv2.resize(screen, (WIDTH,HEIGHT))

            t_minus, t_now, t_plus = prev, prev, prev

            t_now = t_plus
            t_minus = t_now
            t_plus = screen
            t_plus = cv2.blur(t_plus, (4, 4))

            delta_count_last = motion_detection(t_minus, t_now, t_plus)

            prediction = model.predict([screen.reshape(160,120,1)])[0]
            print(prediction)

            mode_choice = np.argmax(prediction)

            turn_thresh = .75
            fwd_thresh = 0.70
#     Choosing the key based on the array
            if mode_choice == 0:
                straight()
                choice_picked = 'straight'
            elif mode_choice == 1:
                reverse()
                choice_picked = 'reverse'
            elif mode_choice == 2:
                left()
                choice_picked = 'left'
            elif mode_choice == 3:
                right()
                choice_picked = 'right'
            elif mode_choice == 4:
                forward_left()
                choice_picked = 'forward+left'
            elif mode_choice == 5:
                forward_right()
                choice_picked = 'forward+right'
            elif mode_choice == 6:
                reverse_left()
                choice_picked = 'reverse+left'
            elif mode_choice == 7:
                reverse_right()
                choice_picked = 'reverse+right'
            elif mode_choice == 8:
                no_keys()
                choice_picked = 'nokeys'

            motion_log.append(delta_count)
            motion_avg = round(mean(motion_log), 3)
            print('loop took {0} seconds. Motion: {1}. Choice: {2}'.format(round(time.time() - last_time, 3), motion_avg, choice_picked))

            if motion_avg < motion_req and len(motion_log) >= log_len:
                print('WERE PROBABLY STUCK FFS, initiating some evasive maneuvers.')

    # 0 = reverse straight, turn left out
    # 1 = reverse straight, turn right out
    # 2 = reverse left, turn right out
    # 3 = reverse right, turn left out

                quick_choice = random.randrange(0, 4)

                if quick_choice == 0:
                    reverse()
                    time.sleep(random.uniform(1, 2))
                    forward_left()
                    time.sleep(random.uniform(1, 2))

                elif quick_choice == 1:
                    reverse()
                    time.sleep(random.uniform(1, 2))
                    forward_right()
                    time.sleep(random.uniform(1, 2))

                elif quick_choice == 2:
                    reverse_left()
                    time.sleep(random.uniform(1, 2))
                    forward_right()
                    time.sleep(random.uniform(1, 2))

                elif quick_choice == 3:
                    reverse_right()
                    time.sleep(random.uniform(1, 2))
                    forward_left()
                    time.sleep(random.uniform(1, 2))

                for i in range(log_len - 2):
                    del motion_log[0]

        keys = key_check()
        # p pauses game and can get annoying.
        if 'T' in keys:
            if paused:
                paused = False
                time.sleep(1)
            else:
                paused = True
                ReleaseKey(A)
                ReleaseKey(W)
                ReleaseKey(D)
                time.sleep(1)

if __name__ == "__main__":
    main()


import torch
import numpy as np
import scipy
from torch.multiprocessing import Queue, RawArray
from ctypes import c_float
import time, datetime
import argparse

from network import Policy
from wrapper import make_env
from util import timing_wrapper, plot
from worker import Worker

class PPOTrainer:
    def __init__(self, args):
        tmp_env = make_env(args.env)
        self.obs_shape = tmp_env.observation_space.shape
        self.num_actions = tmp_env.action_space.n
        self.c_in = self.obs_shape[0]
        del tmp_env

        self.horizon    = args.horizon
        self.eta        = args.eta
        self.epoch      = args.epoch
        self.batch_size = args.batch * args.actors
        self.gamma      = args.gamma
        self.lam        = args.lam
        self.num_actors = args.actors
        self.eps        = args.eps
        self.num_iter   = (args.epoch * args.actors * args.horizon) // self.batch_size # how many times to run SGD on the buffer

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        self.queues = [Queue() for i in range(self.num_actors)]
        self.barrier = Queue()  # This is used as a waiting mechanism, to wait for all the agents to env.step()
        self.score_channel = Queue()

        # these are shmem np.arrays
        self.state, self.reward, self.finished = self.init_shared()

        self.workers = [
            Worker(i, args.env, self.queues[i], self.barrier, self.state, self.reward, self.finished, self.score_channel) for i in range(self.num_actors)
        ]
        self.start_workers()

        self.model  = Policy(self.c_in, self.num_actions).to(self.device)
        self.optim = torch.optim.Adam(self.model.parameters(), lr=self.eta)

        # used for logging and graphing
        self.stat = {
            'scores': [],
            'steps': [],
            'clip_losses': [],
            'value_losses': [],
            'entropies': []
        }

    def init_shared(self):
        state_shape = (self.num_actors, *self.obs_shape)
        scalar_shape = (self.num_actors, 1)

        state = np.empty(state_shape, dtype=np.float32)
        state = RawArray(c_float, state.reshape(-1))
        state = np.frombuffer(state, c_float).reshape(state_shape)

        reward = np.empty(scalar_shape, dtype=np.float32)
        reward = RawArray(c_float, reward.reshape(-1))
        reward = np.frombuffer(reward, c_float).reshape(scalar_shape)

        finished = np.empty(scalar_shape, dtype=np.float32)
        finished = RawArray(c_float, finished.reshape(-1))
        finished = np.frombuffer(finished, c_float).reshape(scalar_shape)

        return state, reward, finished

    def start_workers(self):
        for worker in self.workers:
            worker.start()

    def initialize_state(self):
        for i in range(self.num_actors):
            self.queues[i].put(-1)
        self.wait_for_agents()

    @timing_wrapper
    def broadcast_actions(self, actions):
        actions = actions.cpu().numpy()
        for i in range(self.num_actors):
            self.queues[i].put(actions[i])
        self.wait_for_agents()

        next_state = torch.tensor(self.state).to(self.device)
        reward = torch.tensor(self.reward).to(self.device)
        done = torch.tensor(self.finished).to(self.device)
        return next_state, reward, done

    def wait_for_agents(self):
        for i in range(self.num_actors):
            self.barrier.get()

    def setup_scheduler(self, T_max):
        num_steps = T_max // (self.horizon * self.num_actors)
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optim, lambda x: max(1 - x / num_steps, 0))

    @timing_wrapper
    def train(self, T_max, graph_name=None):
        self.setup_scheduler(T_max)

        global_step = 0

        self.initialize_state()
        state = torch.tensor(self.state).to(self.device)
        while global_step < T_max:

            states      = []
            actions     = []
            rewards     = []
            finished    = []
            sampled_lps = []  # sampled log probabilities
            values      = []

            time_start = time.time()
            duration_fwd = 0
            with torch.no_grad():
                for t in range(self.horizon):
                    global_step += self.num_actors

                    logit, value = self.model(state)
                    prob = torch.softmax(logit, dim=1)
                    log_prob = torch.log_softmax(logit, dim=1)

                    action = prob.multinomial(1)
                    sampled_lp = log_prob.gather(1, action)

                    (next_state, reward, done), duration_brdcst = self.broadcast_actions(action)

                    # appending to buffer
                    states.append(state)
                    actions.append(action)
                    rewards.append(reward)
                    finished.append(done)
                    sampled_lps.append(sampled_lp)
                    values.append(value)

                    state = next_state

                    duration_fwd += duration_brdcst

                _, V = self.model(next_state)
                values.append(V)

            time_forward = time.time()

            # GAE estimation
            GAEs, duration_GAE = self.compute_GAE(rewards, finished, values)

            duration_backward = self.run_gradient_descent(states, actions, sampled_lps, values, GAEs)

            time_end = time.time()

            total_duration = time_end - time_start
            percent_broadcast = duration_fwd / total_duration * 100
            percent_forward = (time_forward - time_start) / total_duration * 100
            percent_GAE = duration_GAE / total_duration * 100
            percent_backward = duration_backward / total_duration * 100

            # print(f"<Time> Total: {total_duration:.2f} | forward: {percent_forward:.2f}% (broadcast {percent_broadcast:.2f}%) | GAE: {percent_GAE:.2f}% | backward: {percent_backward:.2f}%")
            if global_step % (self.num_actors * self.horizon * 30) == 0:
                while not self.score_channel.empty():
                    score, step = self.score_channel.get()
                    self.stat['scores'].append(score)
                    self.stat['steps'].append(step)
                now = datetime.datetime.now().strftime("%H:%M")
                print(f"Step {global_step} | Mean of last 10 scores: {np.mean(self.stat['scores'][-10:]):.2f} | Time: {now}")
                if graph_name is not None:
                    plot(global_step, self.stat, graph_name)
        # Finish
        plot(global_step, self.stat, graph_name)

    @timing_wrapper
    def compute_GAE(self, rewards, finished, values):
        GAEs = []
        advantage = 0
        for i in reversed(range(self.horizon)):
            td_error = rewards[i] + (1 - finished[i]) * self.gamma * values[i + 1] - values[i]
            advantage = td_error + (1 - finished[i]) * self.gamma * self.lam * advantage
            GAEs.append(advantage)
        GAEs = torch.cat(GAEs[::-1]).to(self.device)

        # NOTE: Below is currently not in use because I don't know how to incorporate the 'finished' tensor into account
        # NOTE: This version is much, much faster than the python-looped version above
        # NOTE: But in terms of the total time taken, it doesn't make much of a difference. (~2% compared to ~0.05%)
        # rewards = torch.stack(rewards)
        # finished = torch.stack(finished)
        # values = torch.stack(values)

        # td_error = rewards + (1 - finished) * self.gamma * values[1:] - values[:-1]
        # td_error = td_error.cpu()

        # GAEs = scipy.signal.lfilter([1], [1, -self.gamma * self.lam], td_error.flip(dims=(0,)), axis=0)
        # GAEs = np.flip(GAEs, axis=0)  # flip it back again
        # GAEs = GAEs.reshape(-1, GAEs.shape[-1])  # (horizon, num_actors, 1) --> (horizon * num_actors, 1)
        # GAEs = torch.tensor(GAEs).float().to(self.device)

        return GAEs

    @timing_wrapper
    def run_gradient_descent(self, states, actions, sampled_lps, values, GAEs):

        states      = torch.cat(states)
        actions     = torch.cat(actions)
        sampled_lps = torch.cat(sampled_lps)
        values      = torch.cat(values[:-1])
        targets     = GAEs + values

        self.scheduler.step()
        # Running SGD for K epochs
        for it in range(self.num_iter):
            # Batch indices
            idx = np.random.randint(0, self.horizon * self.num_actors, self.batch_size)

            state      = states[idx]
            action     = actions[idx]
            sampled_lp = sampled_lps[idx]
            GAE        = GAEs[idx]
            value      = values[idx]
            target     = targets[idx]

            # Normalize advantages
            GAE = (GAE - GAE.mean()) / (GAE.std() + 1e-8)

            logit_new, value_new = self.model(state)
            # Clipped values are needed because sometimes values can unexpectedly get really big
            clipped_value_new = value + torch.clamp(value_new - value, -self.eps, self.eps)

            # Calculating policy loss
            prob_new = torch.softmax(logit_new, dim=1)
            lp_new = torch.log_softmax(logit_new, dim=1)
            entropy = -(prob_new * lp_new).sum(1).mean()

            sampled_lp_new = lp_new.gather(1, action)

            ratio = torch.exp(sampled_lp_new - sampled_lp)
            surr1 = ratio * GAE
            surr2 = torch.clamp(ratio, 1 - self.eps, 1 + self.eps) * GAE
            clip_loss = torch.min(surr1, surr2).mean()

            # Calculating value loss
            value_loss1 = (value_new - target).pow(2)
            value_loss2 = (clipped_value_new - target).pow(2)
            value_loss = 0.5 * torch.max(value_loss1, value_loss2).mean()

            final_loss = -clip_loss + value_loss - 0.01 * entropy

            self.optim.zero_grad()
            final_loss.backward()

            # total_norm = 0
            # for p in self.model.parameters():
            #     param_norm = p.grad.data.norm(2)
            #     total_norm += param_norm.item() ** 2
            # total_norm = total_norm ** (1. / 2)
            # print(total_norm)

            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1)
            self.optim.step()

            # graphing
            self.stat['clip_losses'].append(clip_loss.item())
            self.stat['value_losses'].append(value_loss.item())
            self.stat['entropies'].append(entropy.item())

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--env', default='PongNoFrameskip-v4', type=str, help="Name of the environment to use")
    parser.add_argument('--graph', default='graph.png', type=str, help="Name of the graph (must append .png, .jpeg etc)")

    parser.add_argument('--Tmax', default=1e+7, type=float, help="Max number of frames. Default is 10 million")
    parser.add_argument('--horizon', default=128, type=int, help="Horizon (T)")
    parser.add_argument('--eta', default=2.5*1e-4, type=float, help="Learning rate")
    parser.add_argument('--epoch', default=3, type=int, help="K, the number of epochs to run SGD on")
    parser.add_argument('--batch', default=32, type=int, help="Batch size per actor. If batch is K and there are N actors, the total batch size is NK")
    parser.add_argument('--gamma', default=0.99, type=float, help="Gamma (for discount)")
    parser.add_argument('--lam', default=0.95, type=float, help="Lambda (for GAE)")
    parser.add_argument('--actors', default=8, type=int, help="Number of actors")
    parser.add_argument('--eps', default=0.2, type=float, help="Eps, the clipping parameter")
    parser.add_argument('--save', default="", type=str, help="Location to save the model. Default: <None>")

    args = parser.parse_args()

    print(f"""<Configuration>
    Environment:      {args.env}
    Graph:            {args.graph}
    Max frames:       {args.Tmax}
    Horizon (T):      {args.horizon}
    Learning rate:    {args.eta}
    Epoch (optim):    {args.epoch}
    Total Batch size: {args.batch}x{args.actors}={args.batch * args.actors}
    Gamma (discount): {args.gamma}
    Lambda (GAE):     {args.lam}
    Num actors:       {args.actors}
    Epsilon (clip):   {args.eps}
    Current time:     {datetime.datetime.now()}
    Save location:    {args.save if args.save else "-"}
    """)

    trainer = PPOTrainer(args=args)
    time_taken = trainer.train(args.Tmax, args.graph)

    hms = str(datetime.timedelta(seconds=time_taken)).split(':')

    if args.save != "":
        torch.save(trainer.model, args.save)

    print(f"Total time taken: {hms[0]} hours {hms[1]} minutes (={int(time_taken)} seconds)")


import sys, os, glob
import json
import argparse
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
import torch.optim as optim
from torch.optim import lr_scheduler

import smplx
from human_body_prior.tools.model_loader import load_vposer

from utils import GeometryTransformer, BodyParamParser
import time

import os
import pickle

import random
import chamfer_pytorch.dist_chamfer as ext
from torch.autograd import Variable

def cal_loss(xhr, xhr_rec, body, cam_extrinsic, scene_points,
            sdf, s_grid_min_batch, s_grid_max_batch, input_list, iteration, vposer, body_mesh_model_input, body_mesh_model_batch, threshold, contact_id_folder, device):

    contact_part=['back','gluteus','L_Hand','R_Hand','L_Leg','R_Leg','thighs']
    #xhr_rec=torch.cat([xhr_rec,xhr[:,6:]],dim=1)
    ### reconstruction loss
    loss_rec = 0.1*F.l1_loss(xhr[:,:], xhr_rec[:,:])#+0.5*F.l1_loss(xhr[:,6:48], xhr_rec[:,6:48])
    loss_rec = loss_rec#+0.1*F.l1_loss(xhr[xhr_rec.shape[0]//2,:6], xhr_rec[xhr_rec.shape[0]//2,:6])\
                #+0.2*F.l1_loss(input_list[0], xhr_rec[0])+0.2*F.l1_loss(input_list[1], xhr_rec[-1])

    ### vposer loss
    vposer_pose = xhr_rec[:,6:38]
    loss_vposer = torch.mean(vposer_pose**2)

    ### contact loss

    body_param_rec = BodyParamParser.body_params_encapsulate_batch_nobody_hand(xhr_rec)
    #print(xhr_rec.shape)
    #print(body_param_rec['body_pose'].shape)
    body_param_rec['body_pose'] = vposer.decode(body_param_rec['body_pose'],
                                    output_type='aa').view(xhr_rec.shape[0], -1)
    #print(body.shape)
    #print(xhr_rec.shape)
    body_param_rec['betas']=body.repeat(xhr_rec.shape[0],1)
    #print(body_param_rec['body_pose_vp'].shape)

    #temp_out = body_mesh_model()

    smplx_output = body_mesh_model_batch(return_verts=True,
                                        **body_param_rec
                                        )

    body_verts_batch = smplx_output.vertices #[b, 10475,3]
    body_joints_batch = smplx_output.joints
    #print('joints',body_joints_batch[:,7:9,:].shape)
    loss_motion=0

    left_foot_vid, left_foot_fid = GeometryTransformer.get_contact_id(body_segments_folder=contact_id_folder,
                            contact_body_parts=['L_Leg'])
    right_foot_vid, right_foot_fid = GeometryTransformer.get_contact_id(body_segments_folder=contact_id_folder,
                            contact_body_parts=['R_Leg'])
    gluteus_vid, gluteus_fid = GeometryTransformer.get_contact_id(body_segments_folder=contact_id_folder,
                            contact_body_parts=['gluteus'])

    left_foot_verts_batch = body_verts_batch[:, left_foot_vid, :]
    right_foot_verts_batch = body_verts_batch[:, right_foot_vid, :]
    butt_verts_batch = body_verts_batch[:,gluteus_vid,:]

    #print(body_verts_batch.shape)
    if iteration<20:
        foot_motion=False
    else:
        foot_motion=False

    loss_skating=0

    #temp_list_=[]
    #temp_list2=[]

    #global left_foot_scatter
    #left_foot_scatter=[]
    #global right_foot_scatter
    #right_foot_scatter=[]

    if iteration==0:
        global highest_joints
        highest_joints=torch.min(body_joints_batch[:,7:9,1]).detach()
        global temp_list
        temp_list=[]

        #global foot_list
        foot_list=[]
        for iii in range(1,body_verts_batch.shape[0]-1,1):

            '''
            butt_verts_avg=torch.tensor([torch.mean(butt_verts_batch[iii][:,:1]),torch.mean(butt_verts_batch[iii][:,1:2]),
                                            torch.mean(butt_verts_batch[iii][:,2:3])])
            left_foot_avg=torch.tensor([torch.mean(left_foot_verts_batch[iii][:,:1]),torch.mean(left_foot_verts_batch[iii][:,1:2]),
                                            torch.mean(left_foot_verts_batch[iii][:,2:3])])
            right_foot_avg=torch.tensor([torch.mean(right_foot_verts_batch[iii][:,:1]),torch.mean(right_foot_verts_batch[iii][:,1:2]),
                                            torch.mean(right_foot_verts_batch[iii][:,2:3])])

            next_left_foot_avg=torch.tensor([torch.mean(left_foot_verts_batch[iii+1][:,:1]),torch.mean(left_foot_verts_batch[iii+1][:,1:2]),
                                            torch.mean(left_foot_verts_batch[iii+1][:,2:3])])

            next_right_foot_avg=torch.tensor([torch.mean(right_foot_verts_batch[iii+1][:,:1]),torch.mean(right_foot_verts_batch[iii+1][:,1:2]),
                                            torch.mean(right_foot_verts_batch[iii+1][:,2:3])])
            '''
            left_foot_joint=body_joints_batch[iii,7,:]
            right_foot_joint=body_joints_batch[iii,8,:]

            left_knee_joint=body_joints_batch[iii,4,:]
            right_knee_joint=body_joints_batch[iii,5,:]

            #left_foot_scatter.append(left_foot_joint.detach().numpy())
            #right_foot_scatter.append(right_foot_joint.detach().numpy())

            next_left_foot_joint=body_joints_batch[iii+1,7,:]
            next_right_foot_joint=body_joints_batch[iii+1,8,:]

            last_left_foot_joint=body_joints_batch[iii-1,7,:]
            last_right_foot_joint=body_joints_batch[iii-1,8,:]

            left_joint_distance=torch.mean((next_left_foot_joint-left_foot_joint)**2)
            right_joint_distance=torch.mean((next_right_foot_joint-right_foot_joint)**2)

            left_joint_longer_distance=torch.mean((next_left_foot_joint-last_left_foot_joint)**2)
            right_joint_longer_distance=torch.mean((next_right_foot_joint-last_right_foot_joint)**2)

            #left_avg_low=left_foot_avg[1:2]
            #right_avg_low=right_foot_avg[1:2]
            #if left_avg_distance>right_avg_distance:
            #    temp_list.append('left')
                #foot_list.append(left_foot_verts_batch[iii])
            #else:
            #    temp_list.append('right')
                #foot_list.append(right_foot_verts_batch[iii])
            #threshold=threshold
            if threshold*right_joint_distance<left_joint_distance and threshold*right_joint_longer_distance<left_joint_longer_distance:
            #if body_joints_batch[iii,5,1]>1.1*body_joints_batch[iii,4,1]:
                temp_list.append('right')
                #foot_list.append(right_foot_verts_batch[iii])
            elif right_joint_distance>threshold*left_joint_distance and right_joint_longer_distance>threshold*left_joint_longer_distance:
            #elif 1.1*body_joints_batch[iii,5,1]<body_joints_batch[iii,4,1]:
                temp_list.append('left')
                #foot_list.append(left_foot_verts_batch[iii])
            else:
                temp_list.append('unknown')
                #foot_list.append(None)
#        #print(temp_list)
        for jj in range(len(temp_list)):
            if jj==0 and temp_list[jj]=='unknown':
                for jjj in range(jj+1,len(temp_list)):
                    if temp_list[jjj]!='unknown':
                        for jjjj in range(jj,jjj):
                            temp_list[jjjj]=temp_list[jjj]
                        break
            if jj==len(temp_list)-1 and temp_list[jj]=='unknown':
                for jjj in range(jj-1,0,-1):
                    if temp_list[jjj]!='unknown':
                        for jjjj in range(jj,jjj,-1):
                            temp_list[jjjj]=temp_list[jjj]
                        break
            if temp_list[jj]!='unknown':
                for jjj in range(jj+1,len(temp_list)):

                    if temp_list[jjj]=='unknown':
                        continue
                    elif temp_list[jjj]==temp_list[jj]:
                        for jjjj in range(jj,jjj+1):
                            temp_list[jjjj]=temp_list[jj]
                        break
                    else:
                        for jjjj in range(jj,jjj+1):
                            #temp_list[jjjj]=temp_list[jj]
                            if jjjj<int((jj+jjj)/2):
                                temp_list[jjjj]=temp_list[jj]
                            else:
                                temp_list[jjjj]=temp_list[jjj]
                        break
        temp_list.insert(0,temp_list[0])
        temp_list.append(temp_list[-1])
        print(temp_list)

        for kkk in range(len(temp_list)):
            if temp_list[kkk] == 'left':
                foot_list.append(left_foot_verts_batch[kkk])
            if temp_list[kkk] =='right':
                foot_list.append(right_foot_verts_batch[kkk])
        #print('foot_list',foot_list.shape)
        global temp_list2
        temp_list2=[]
        global avg_foot_list
        avg_foot_list=[]
        start=0
        for jj in range(len(temp_list)-1):
            if temp_list[jj]==temp_list[jj+1]:

                pass#continue
            if temp_list[jj]!=temp_list[jj+1]:
                temp_list2.append([jj,temp_list[jj]])
                end=jj
                part_avg_foot_list=foot_list[start:end+1]
                part_avg_foot=sum(part_avg_foot_list)/(len(part_avg_foot_list))

                for each in range(len(part_avg_foot_list)):
                    avg_foot_list.append(part_avg_foot.detach())
                start=jj+1
        if start<len(temp_list):
            part_avg_foot_list=foot_list[start:]
            part_avg_foot=sum(part_avg_foot_list)/(len(part_avg_foot_list))
            #print(part_avg_foot.shape)

            for each in range(len(part_avg_foot_list)):
                    avg_foot_list.append(part_avg_foot.detach())
        global midlist
        midlist=[]
        for q in range(len(temp_list2)):
            if q==0:
                temp_highest=torch.min(body_joints_batch[:temp_list2[q][0],7:9,1]).detach()
                midlist.append([int(0.5*temp_list2[q][0]),temp_list2[q][1],temp_highest])
                if len(temp_list2)==1:
                    temp_highest=torch.min(body_joints_batch[temp_list2[q][0]:,7:9,1]).detach()
                    midlist.append([int(0.5*(temp_list2[q][0]+len(avg_foot_list))),temp_list[temp_list2[q][0]+1],temp_highest])

                else:
                    temp_highest=torch.min(body_joints_batch[temp_list2[q][0]:temp_list2[q+1][0],7:9,1]).detach()
                    midlist.append([int(0.5*(temp_list2[q][0]+temp_list2[q+1][0])),temp_list2[q+1][1],temp_highest])

            elif q==len(temp_list2)-1:
                temp_highest=torch.min(body_joints_batch[temp_list2[q][0]:,7:9,1]).detach()
                midlist.append([int(0.5*(temp_list2[q][0]+len(avg_foot_list))),temp_list2[q-1][1],temp_highest])
            else:
                temp_highest=torch.min(body_joints_batch[temp_list2[q][0]:temp_list2[q+1][0],7:9,1]).detach()
                midlist.append([int(0.5*(temp_list2[q][0]+temp_list2[q+1][0])),temp_list2[q+1][1],temp_highest])

    #print(temp_list)
    #print(len(temp_list))
    #print(temp_list2)
    #temp_list_.append('unknown')
    #for i in range(1,len(temp_list)-1,1):
    #    if temp_list[i]==temp_list[i-1] and temp_list[i]==temp_list[i+1]:
    #        temp_list_.append(temp_list[i])
    #    else:
    #        temp_list_.append('unknown')
    #print(temp_list_)
    #print(foot_list)
    '''
    if iteration==0:
        global temp_list2
        temp_list2=[]
        global avg_foot_list
        avg_foot_list=[]
        start=0
        for jj in range(len(temp_list)-1):
            if temp_list[jj]==temp_list[jj+1]:

                pass#continue
            if temp_list[jj]!=temp_list[jj+1]:
                temp_list2.append([jj,temp_list[jj]])
                end=jj
                part_avg_foot_list=foot_list[start:end+1]
                part_avg_foot=sum(part_avg_foot_list)/(len(part_avg_foot_list))

                for each in range(len(part_avg_foot_list)):
                    avg_foot_list.append(part_avg_foot)
                start=jj+1
        if start<len(temp_list):
            part_avg_foot_list=foot_list[start:]
            part_avg_foot=sum(part_avg_foot_list)/(len(part_avg_foot_list))
            #print(part_avg_foot.shape)

            for each in range(len(part_avg_foot_list)):
                    avg_foot_list.append(part_avg_foot)
    #print(len(avg_foot_list))
    #print(len(foot_list))
    '''

    #print(avg_foot_list[0])
    foot_list=[]

    for kkk in range(len(temp_list)):
        if temp_list[kkk] == 'left':
            foot_list.append(left_foot_verts_batch[kkk])
        if temp_list[kkk] =='right':
            foot_list.append(right_foot_verts_batch[kkk])
#    print(loss_skating)

    for kk in range(len(foot_list)):
        temp_loss=torch.mean((foot_list[kk]-avg_foot_list[kk])**2)
        loss_skating=loss_skating+temp_loss
#    print(loss_skating)
#    print('temp_list2:',temp_list2)
    #print(avg_foot_list)
    for qq in range(len(temp_list2)):
        each=temp_list2[qq]
        if each[0]<len(right_foot_verts_batch)-1:
            if each[1]=='left':
                #temp_loss_left=torch.mean((avg_foot_list[each[0]]-left_foot_verts_batch[each[0]+1])**2)\
                #           +torch.mean((avg_foot_list[each[0]]-left_foot_verts_batch[each[0]])**2)

                temp_loss_=torch.mean((right_foot_verts_batch[each[0]]-avg_foot_list[each[0]+1])**2)
                                        #torch.mean((right_foot_verts_batch[each[0]]-avg_foot_list[each[0]+1])**2)

            if each[1]=='right':
                temp_loss_=torch.mean((left_foot_verts_batch[each[0]]-avg_foot_list[each[0]+1])**2)
                #temp_loss_right=torch.mean((avg_foot_list[each[0]]-right_foot_verts_batch[each[0]+1])**2)\
                #                    +torch.mean((avg_foot_list[each[0]]-right_foot_verts_batch[each[0]])**2)

            #print(temp_loss_left)
            #print(temp_loss_)
            loss_skating=loss_skating+temp_loss_
    loss_height=0
    for each in midlist:
        if each[0]<len(right_foot_verts_batch)-1:
            if each[1]=='left':

                loss_height=loss_height+(body_joints_batch[each[0],8,1]-each[2])**2
            if each[1]=='right':

                loss_height=loss_height+(body_joints_batch[each[0],7,1]-each[2])**2
#    print('height:',loss_height)
    loss_skating=loss_skating+0.05*loss_height

    #loss_contact

    body_param_rec_input = BodyParamParser.body_params_encapsulate_batch_nobody_hand(input_list)
    #print('2:',body_param_rec_input['body_pose_vp'].shape)

    body_param_rec_input['body_pose'] = vposer.decode(body_param_rec_input['body_pose'], output_type='aa').view(input_list.shape[0], -1)
    #print('3:',joint_rot_batch.shape)
    body_param_rec_input['betas']=body.repeat(1,1)

    #body_param_rec_input['betas']=body.repeat(input_list.shape[0])
    smplx_input = body_mesh_model_input(return_verts=True,
                                    **body_param_rec_input
                                    )
    body_verts_batch_input = smplx_input.vertices
    #print('4:',body_verts_batch_input.shape)
    #print()
    left_foot_verts_batch_input = body_verts_batch_input[:, left_foot_vid, :]
    right_foot_verts_batch_input = body_verts_batch_input[:, right_foot_vid, :]

    left_start=torch.mean((left_foot_verts_batch[0]-left_foot_verts_batch_input[0])**2)
    right_start=torch.mean((right_foot_verts_batch[0]-right_foot_verts_batch_input[0])**2)

    left_end=torch.mean((left_foot_verts_batch[-1]-left_foot_verts_batch_input[1])**2)
    right_end=torch.mean((right_foot_verts_batch[-1]-right_foot_verts_batch_input[1])**2)

    for iii in range(body_verts_batch.shape[0]-1):
        '''
        butt_verts_avg=torch.tensor([torch.mean(butt_verts_batch[iii][:,:1]),torch.mean(butt_verts_batch[iii][:,1:2]),
                                        torch.mean(butt_verts_batch[iii][:,2:3])])
        left_foot_avg=torch.tensor([torch.mean(left_foot_verts_batch[iii][:,:1]),torch.mean(left_foot_verts_batch[iii][:,1:2]),
                                        torch.mean(left_foot_verts_batch[iii][:,2:3])])
        right_foot_avg=torch.tensor([torch.mean(right_foot_verts_batch[iii][:,:1]),torch.mean(right_foot_verts_batch[iii][:,1:2]),
                                        torch.mean(right_foot_verts_batch[iii][:,2:3])])
        '''
        left_foot_loss=torch.mean((left_foot_verts_batch[iii]-left_foot_verts_batch[iii+1])**2)
        right_foot_loss=torch.mean((right_foot_verts_batch[iii]-right_foot_verts_batch[iii+1])**2)
        loss_motion=loss_motion+(left_foot_loss+right_foot_loss)

        loss_motion=loss_motion+torch.mean((body_verts_batch[iii]-body_verts_batch[iii+1])**2)

    loss_start_end=torch.mean((body_verts_batch[0]-body_verts_batch_input[0])**2)+torch.mean((body_verts_batch[-1]-body_verts_batch_input[1])**2)
    loss_motion=loss_motion+loss_start_end+min(left_end,right_end)+min(left_start,right_start)
    #/+5*min(left_start,right_start)+5*min(left_end,right_end)
    loss_contact=0

    #body_verts_batch__ = GeometryTransformer.verts_transform(body_verts_batch, torch.tensor(cam_extrinsic,dtype=torch.float32))#.cuda())
    body_verts_batch__ = GeometryTransformer.verts_transform(body_verts_batch, cam_extrinsic)
    vid, fid = GeometryTransformer.get_contact_id(body_segments_folder=contact_id_folder,
                            contact_body_parts=contact_part)
    foot_vid, foot_fid = GeometryTransformer.get_contact_id(body_segments_folder=contact_id_folder,
                            contact_body_parts=['L_Leg','R_Leg'])

    body_verts_contact_batch = body_verts_batch[:, vid, :]

    dist_chamfer_contact = ext.chamferDist()

    contact_dist, _ = dist_chamfer_contact(
                                    body_verts_contact_batch.contiguous(),
                                    scene_points.contiguous()
                                    )

    loss_contact = (1 *
                            torch.mean( torch.sqrt(contact_dist+1e-4)
                            /(torch.sqrt(contact_dist+1e-4)+1.0)  )  )

    #loss_sdf

    loss_sdf_pene=0

    for ttt in range(body_verts_batch__.shape[0]):

        body_verts_batch_=body_verts_batch__[ttt:ttt+1,:,:]

        norm_verts_batch = ((body_verts_batch_ - s_grid_min_batch.unsqueeze(1))
                                / (s_grid_max_batch.unsqueeze(1) - s_grid_min_batch.unsqueeze(1)) *2 -1)

        n_verts = norm_verts_batch.shape[1]

        body_sdf_batch = F.grid_sample(sdf.unsqueeze(1),
                        norm_verts_batch[:,:,[2,1,0]].view(-1, n_verts,1,1,3),
                        padding_mode='border')

        # if there are no penetrating vertices then set sdf_penetration_loss = 0
        if body_sdf_batch.lt(0).sum().item() < 1:
            loss_sdf_pene_ = torch.tensor(0.0, dtype=torch.float32,device=device)

        else:
            loss_sdf_pene_ = body_sdf_batch[body_sdf_batch < 0].abs().mean()
        loss_sdf_pene=loss_sdf_pene+loss_sdf_pene_
        #print(loss_sdf_pene)

    #loss_sdf_pene=0

    return loss_rec, loss_vposer, loss_contact, loss_sdf_pene, loss_motion, loss_skating

def fitting(xhr_in, body, cam_extrinsic,
            s_verts, s_sdf, s_grid_min, s_grid_max,
            fittingconfig, lossconfig, start_end, vposer, body_mesh_model_batch, threshold=1.6, contact_id_folder='./data/body_segments', device='cuda'):

    body_mesh_model_input = smplx.create('./models',
                               model_type='smplx',
                               gender='neutral', ext='npz',
                               num_pca_comps=12,
                               create_global_orient=True,
                               create_body_pose=True,
                               create_betas=True,
                               create_left_hand_pose=True,
                               create_right_hand_pose=True,
                               create_expression=True,
                               create_jaw_pose=True,
                               create_leye_pose=True,
                               create_reye_pose=True,
                               create_transl=True,
                               batch_size=2
                               )

    body_mesh_model_input=body_mesh_model_input.to(device)

    batch_size = xhr_in.shape[0]

    v = Variable(torch.randn(batch_size,62).to(device)
                , requires_grad=True)

    optimizer = optim.Adam([v], lr=fittingconfig['init_lr_h'])
    v.data = xhr_in.clone()

    cam_ext_batch = cam_extrinsic.repeat(batch_size, 1,1)
    #max_d_batch = max_d.repeat(batch_size)
    s_verts_batch = s_verts.repeat(batch_size, 1,1)
    s_sdf_batch = s_sdf.repeat(1, 1,1,1)
    s_grid_min_batch = s_grid_min.repeat(1, 1)
    s_grid_max_batch = s_grid_max.repeat(1, 1)

    for ii in range(fittingconfig['num_iter']):

        optimizer.zero_grad()

        loss_rec, loss_vposer, loss_contact, loss_collision, loss_motion, loss_skating = cal_loss(xhr_in, v, body, cam_ext_batch, s_verts_batch,
                                                                        s_sdf_batch,s_grid_min_batch, s_grid_max_batch,start_end
                                                                        , ii, vposer, body_mesh_model_input, body_mesh_model_batch, threshold, contact_id_folder, device)

        loss_rec=lossconfig['weight_loss_rec']*loss_rec
        loss_contact=lossconfig['weight_contact']*loss_contact
        loss_collision=lossconfig['weight_collision']*loss_collision

        loss_motion=lossconfig['weight_motion']*loss_motion
        loss_skating=lossconfig['weight_skating']*loss_skating

        loss = loss_rec + loss_contact + loss_collision+loss_motion+loss_skating
        '''
        if fittingconfig['verbose']:
            print('[INFO][fitting] iter={:d}, l_rec={:f},loss_skating={:f} ,l_contact={:f}, l_collision={:f},l_motion={:f},l_vposer={:f}'.format(
                                    ii, loss_rec, loss_skating,
                                    loss_contact, loss_collision,loss_motion,loss_vposer))
        '''
        loss.backward(retain_graph=True)
        optimizer.step()

    ### recover global translation and orientation
    #xh_rec = convert_to_3D_rot(xhr_rec)
    #xhr_rec=torch.cat([v,xhr_in[:,6:]],dim=1)
    return v


#! /usr/bin/env python

'''
Create a mirage-format seed image from
an input fits file. This fits file must contain a
valid WCS, and be distortion-free.

This code will extract a sub-image from the input
fits file. The subimage is centered at the input
RA and Dec (crop_center_ra, crop_center_dec) and
its size is determined by the aperture input.

The extracted sub-image is then blotted onto the
requested NIRCam detector's FOV, including the
distortion model.

The resulting image is then saved in such a format
that it can be used as input in the obs_generator.py,
image_simulator.py, or disperser code in the
mirage package. Essentially this script
is designed to take the place of catalog_seed_image.py
in the case where the user has a distortion-free
mosaic image that they wish to break up into
NIRCam images.

Another way to use this code would be for an input
fits image of some extended object that the user
wishes to distort according to the NIRCam distortion
model. This distorted file can then be placed in
an extended object catalog which will be used by
the catalog_seed_image.py step. In this method,
multiple extended objets could be distorted and
added to the same frame.

Inputs:
mosaicfile: The name of the fits file containing the
            distortion-free image to be blotted. The code
            currently assumes that the mosaicfile is
            oriented north-up.

From here, there are two ways to call fits_seed_image.py

1. parameter file: a yaml input file matching the format
   (optional)      required by other steps of the nircam
                   data simulator

2. Manual inputs:
   aperture -
              aperture name matching an entry in the
              subarray definition file (e.g. NRCB5_FULL)

   crop_center_ra, crop_center_dec -
              The RA and Dec at the center of the sub-image
              to crop from mosaicfile. This, in combination
              with the array size (determined from aperture)
              define the sub-image to use.

   blot_center_ra, blot_center_dec -
              The RA and Dec at the center of the blotted
              sub-image. If these are equal to crop_center_ra
              and dec, then the center of the cropped image
              will remain at the center of the blotted image.

   blot_pav3 -
              Position angle of the blotted image.

   subarray_defs -
              Name of the ascii file containing the definitions
              of the possible subarrays. By setting this value
              equal to 'config', the code will use the
              subarray definition file packaged with the
              NIRCam data simulator.

   flux_cal_file -
              Ascii file listing the zeropoints for all
              NIRCam filters. By setting this value equal
              to 'config', the code will use the flux
              calibration file that is packaged with the
              NIRCam data simulator. This information is
              used only to populate the PHOTFLAM, PHOTFNU,
              and PHOTPLAM header keywords.

   filter -
              The NIRCam filter to use. This is used to
              select the appropriate row from the flux
              calibration file.

   pupil -
              The NIRCam pupil value to use. This is used
              to select the appropriate row from the flux
              calibration file in the case where a filter
              that is located in the pupil wheel is used.
              If the pupil value does not start with 'F'
              (e.g. F405N), then it is ignored.

   grism_source_image -
              True/False. If you intend to send the seed
              image output to the disperser software, then
              set this value to True. The output image will
              then be made larger than the nominal aperture
              size by a factor of sqrt(2), in order for the
              disperser to keep track of sources just outside
              the nominal FOV, whose dispersed spectra may
              fall onto the detector.

   outfile -
              Name of the file to contain the output seed
              image.

   outdir -
              Directory in which to place outfile

Outputs:
self.seedimage, self.seed_segmap, self.seedinfo
contain the seed (countrate) image, the associated
segmentation map, and header information required
by subsequent steps of the nircam data simulator.
The segmentation map is only used in the case where
the seed image is dispersed by a call to the
disperser software.

Example calls:

1. Using parameter file
s = fits_seed_image.ImgSeed('fits_seed_image_test.yaml')
s.mosaicfile = 'large_distortionFree_mosaic_drz.fits'
s.crop_and_blot()

2. Manual inputs
s = fits_seed_image.ImgSeed()
s.mosaicfile = 'large_distortionFree_mosaic_drz.fits'
s.aperture = 'NRCB5_FULL'
s.crop_center_ra = 53.1
s.crop_center_dec = -27.8
s.blot_center_ra = 53.1
s.blot_center_de = -27.8
s.blot_pav3 = 0.
s.subarray_defs = 'config'
s.flux_cal_file = 'config'
s.filter = 'F250M'
s.pupil = 'CLEAR'
s.grism_source_image = False
s.outfile = 'test_mosaic_seed.fits'
s.outdir = '/my/simulation/outputs/'
s.crop_and_blot()
'''

import os
import sys
import argparse

import yaml
import pkg_resources
import numpy as np
from math import radians
from astropy.io import fits, ascii

from . import crop_mosaic, blot_image

config_files = {'subarray_defs':'NIRCam_subarray_definitions.list',
                'flux_cal':'NIRCam_zeropoints.list'}

class ImgSeed:
    def __init__(self,paramfile = None):
        self.mosaicfile = None
        self.crop_center_ra = 0.
        self.crop_center_dec = 0.
        self.channel = None
        self.detector = None
        self.blot_center_ra = 0.
        self.blot_center_dec = 0.
        self.blot_pav3 = 0.
        self.aperture = ''
        self.subarray_defs = ''
        self.flux_cal_file = ''
        self.filter = ''
        self.pupil = ''
        self.grism_source_image = False
        self.outfile = 'seed_image_from_mosaic.fits'
        self.outdir = './'
        self.grism_direct_factor = np.sqrt(2.)

        # Locate the module files, so that we know where to look
        # for config subdirectory
        self.modpath = pkg_resources.resource_filename('mirage','')

        # self.coords contains the factor by which the
        # nominal output array size needs to be increased
        # (used for WFSS mode), as well as the coordinate
        # offset between the nominal output array coordinates,
        # and those of the expanded array. These are needed
        # mostly for WFSS observations, where the nominal output
        # array will not sit centered in the expanded output image.
        self.coords = {'x':1.,'xoffset':0.,'y':1.,'yoffset':0.}

        # If a paramfile is provided, read in and
        # set params
        if paramfile is not None:
            self.read_param_file(paramfile)

    def read_param_file(self, file):
        '''If an input yaml file is given
        which matches the format of the input
        yaml files for the other mirage
        steps, read in and set needed parameters'''

        if os.path.isfile(file):
            with open(file,'r') as f:
                params = yaml.safe_load(f)
        else:
            print(("WARNING: {} does not exist."
                   .format(file)))
            sys.exit()

        self.aperture = params['Readout']['array_name']
        self.detector = self.aperture[3:5]
        if '5' in self.detector:
            self.channel = 'long'
        else:
            self.channel = 'short'
        #module = self.detector[0].upper()
        self.crop_center_ra = params['Telescope']['ra']
        self.crop_center_dec = params['Telescope']['dec']
        self.blot_center_ra = params['Telescope']['ra']
        self.blot_center_dec = params['Telescope']['dec']
        self.blot_pav3 = params['Telescope']['rotation']
        self.subarray_defs = params['Reffiles']['subarray_defs']
        self.flux_cal_file = params['Reffiles']['flux_cal']
        self.filter = params['Readout']['filter']
        self.pupil = params['Readout']['pupil']
        # If input is to be dispersed later, we need to
        # expand the region to be blotted
        if params['Output']['grism_source_image']:
            self.grism_coord_adjust()
        self.outfile = params['Output']['file']
        self.outdir = params['Output']['directory']

    def expand_config_paths(self,file,filetype):
        if file.lower() == 'config':
            sfile = config_files[filetype]
            return os.path.join(self.modpath, 'config', sfile)

    def fluxcal_info(self, fluxfile, filter, module):
        # Read in list of zeropoints/photflam/photfnu
        self.zps = ascii.read(fluxfile)
        mtch = ((self.zps['Filter'] == filter) \
                & (self.zps['Module'] == module))
        self.photflam = self.zps['PHOTFLAM'][mtch][0]
        self.photfnu = self.zps['PHOTFNU'][mtch][0]
        self.pivot = self.zps['Pivot_wave'][mtch][0]

    def aperture_info(self, apname):
        '''
        Return basic information on the reuqested
        aperture. This information comes from the
        subarray defintion file in the config
        directory
        '''
        # Read in the subarray definition file
        print(self.subarray_defs)
        self.read_subarray_definition_file(self.subarray_defs)

        # Find the matching entry in the definition file
        if self.aperture in self.subdict['AperName']:
            mtch = self.aperture == self.subdict['AperName']
            xstart = self.subdict['xstart'].data[mtch][0]
            ystart = self.subdict['ystart'].data[mtch][0]
            xend = self.subdict['xend'].data[mtch][0]
            yend = self.subdict['yend'].data[mtch][0]
            self.array = {'xstart':xstart,
                               'ystart':ystart,
                               'xend':xend,
                               'yend':yend}
        else:
            print(("WARNING: Array name {} does not "
                   .format(apname)))
            print("to be present in subarray definition file.")
            sys.exit()

    def read_subarray_definition_file(self, file):
        '''Read in the file that contains a list
        of subarray names and positions on the detector
        '''
        try:
            self.subdict = ascii.read(file,data_start=1,header_start=0)
        except:
            print("Error: could not read in subarray definitions file:")
            print(file)
            sys.exit()

    def crop_and_blot(self):
        '''Extract an area of the input fits
        file that is slightly larger than the
        FOV, and use the JWST pipeline version
        of blot to resample to the correct
        pixel scale and introduce the proper
        distortion'''

        # Locations of conifg files
        self.subarray_defs = self.expand_config_paths(self.subarray_defs, 'subarray_defs')
        self.flux_cal_file = self.expand_config_paths(self.flux_cal_file, 'flux_cal')

        if self.detector is None:
            self.detector = self.aperture[3:5]
        if self.channel is None:
            if '5' in self.detector:
                self.channel = 'long'
            else:
                self.channel = 'short'

        # Get information on the aperture
        self.aperture_info(self.aperture)
        xstart = self.array['xstart']
        ystart = self.array['ystart']
        xdim = self.array['xend'] - xstart + 1
        ydim = self.array['yend'] - ystart + 1
        module = self.detector[0].upper()

        # Flux calibration information
        if self.pupil[0].upper() == 'F':
            usefilt = self.pupil.upper()
        else:
            usefilt = self.filter.upper()
        self.fluxcal_info(self.flux_cal_file, usefilt, module)

        # Extract
        crop = crop_mosaic.Extraction()
        crop.mosaicfile = self.mosaicfile
        crop.center_ra = self.crop_center_ra
        crop.center_dec = self.crop_center_dec
        crop.channel = self.channel
        crop.corner = (xstart, ystart)
        crop.dimensions = (xdim, ydim)
        crop.extract()

        # Blot
        blot = blot_image.Blot()
        blot.blotfile = crop.cropped
        blot.detector = [self.detector]
        blot.center_ra = [self.blot_center_ra]
        blot.center_dec = [self.blot_center_dec]
        blot.pav3 = [self.blot_pav3]
        blot.blot()

        for dm in blot.blotted_datamodels:
            # Create segmentation map
            # (used only when dispsersing image for WFSS)
            self.seed_segmap = self.make_segmap(dm)

            # Save
            self.save_seed_image(dm)

    def make_segmap(self, model):
        '''Create a segmentation map of the
        input image'''
        noise = np.median(model.data)
        map = detect_sources(model.data,noise*3.,8).data
        return map

    def save_seed_image(self, model):
        '''Save the seed image in the proper
        format that it can be used in subsequent
        steps of mirage'''
        self.seed_image = model.data
        dim = len(self.seed_image.shape)
        if dim == 2:
            units = 'e-/sec'
            yd,xd = arrayshape
            tgroup = 0.
        else:
            print(("WARNING: {}D seed images from "
                   "input fits file not yet supported."
                   .format(dim)))
            sys.exit()

        dot = self.outfile.rfind('.fits')
        if dot == -1:
            dot = len(self.outfile)
        self.basename = os.path.join(self.outdir,
                                     self.outfile[0:dot])
        self.seed_file = self.basename + '_' \
                      + self.filter + '_seed_image.fits'

        xcent_fov = xd / 2
        ycent_fov = yd / 2
        kw = {}
        kw['xcenter'] = xcent_fov
        kw['ycenter'] = ycent_fov
        kw['units'] = units
        kw['TGROUP'] = tgroup
        if self.pupil.upper() == 'F':
            kw['filter'] = self.pupil
        else:
            kw['filter'] = self.filter
        kw['PHOTFLAM'] = self.photflam
        kw['PHOTFNU'] = self.photfnu
        kw['PHOTPLAM'] = self.pivot * 1.e4 #put into angstroms
        kw['NOMXDIM'] = self.array['xend'] - self.array['xstart'] + 1
        kw['NOMYDIM'] = self.array['yend'] - self.array['ystart'] + 1
        kw['NOMXSTRT'] = self.coords['xoffset'] + 1
        kw['NOMXEND'] = self.coords['xoffset'] + kw['NOMXDIM']
        kw['NOMYSTRT'] = self.coords['yoffset'] + 1
        kw['NOMYEND'] = self.coords['yoffset'] + kw['NOMYDIM']
        kw['GRISMPAD'] = self.grism_direct_factor
        self.seedinfo = kw
        self.save_fits(self.seedimage, self.seed_file, key_dict = kw,
                       image2 = self.seed_segmap, image2type = 'SEGMAP')

    def saveSingleFits(self, image, name, key_dict = None,
                       image2 = None, image2type = None):
        # Save an array into the first extension of a fits file
        h0 = fits.PrimaryHDU()
        h1 = fits.ImageHDU(image,name='DATA')
        if image2 is not None:
            h2 = fits.ImageHDU(image2)
            if image2type is not None:
                h2.header['EXTNAME'] = image2type

        # If a keyword dictionary is provided, put the
        # keywords into the 0th and 1st extension headers
        if key_dict is not None:
            for key in key_dict:
                h0.header[key] = key_dict[key]
                h1.header[key] = key_dict[key]

        if image2 is None:
            hdulist = fits.HDUList([h0,h1])
        else:
            hdulist = fits.HDUList([h0,h1,h2])
        hdulist.writeto(name,overwrite=True)

    def grism_coord_adjust(self):
        # Calculate the factors by which to expand the
        # output array size, as well as the coordinate
        # offsets between the nominal output array and
        # the input lists if the observation being
        # modeled is to be dispersed with the grism
        dtor = radians(1.)

        # Normal imaging with grism image requested
        self.coords['x'] = self.grism_direct_factor
        self.coords['y'] = self.grism_direct_factor
        self.coords['xoffset'] = np.int((self.grism_direct_factor - 1.)
                                        * (self.subarray_bounds[2] -
                                           self.subarray_bounds[0]+1) / 2.)
        self.coords['yoffset'] = np.int((self.grism_direct_factor - 1.)
                                        * (self.subarray_bounds[3] -
                                           self.subarray_bounds[1]+1) / 2.)

    def add_options(self, parser = None, usage = None):
        if parser is None:
            parser = argparse.ArgumentParser(usage=usage,\
                        description='Create seed image via catalogs')
        parser.add_argument("paramfile",help='File describing the input parameters and instrument settings to use. (YAML format).')
        parser.add_argument("mosaicfile",help="Fits file containing the image to create the seed image from")
        parser.add_argument("crop_center_ra",help="RA of center of the image to crop from mosaicimage")
        parser.add_argument("crop_center_dec",help="Dec of center of the image to crop from mosaicimage")
        parser.add_argument("channel",help="NIRCam channel: 'short' or 'long'")
        parser.add_argument("detector",help="NIRCam detector: e.g. 'A1','B5'")
        parser.add_argument("blot_center_ra",help="RA of center of the output blotted image")
        parser.add_argument("blot_center_dec",help="Dec of center of the output blotted image")
        parser.add_argument("blot_pav3",help="Position angle of the output blotted image")
        parser.add_argument("aperture",help="Aperture to use. (e.g. NRCA1_FULL)")
        parser.add_argument("subarray_defs",help="Ascii file containing subarray definitions")
        return parser

if __name__ == '__main__':

    usagestring = 'USAGE: fits_seed_image.py inputs.yaml'

    seed = Imgseed()
    parser = seed.add_options(usage = usagestring)
    args = parser.parse_args(namespace=seed)
    seed.read_param_file(args.paramfile)
    seed.crop_and_blot()