#!/usr/bin/env python3
import os, sys 
sys.path.append('/auto')
from   pet_misc  import *

dyn_range  = 0.0001
est_range  = 0.001
flux_main  = 0.0
flux_comp  = 0.0
dist_apr   = 37.0

tot_flux      = 0.0
tot_flux_main = 0.0
tot_flux_comp = 0.0

size  =  2.0
ratio =  0.3
angle = 20.0


if ( len(sys.argv)-1 < 1 ):
     print ( "Usage: gen_inpit_model uvs_file" )
     exit  ( 1 )
else:
     uvs_file    = sys.argv[1]

modin_file  = uvs_file.replace("_uvs.fits",".mod")
modout_file = uvs_file.replace("uvs.fits","fit.mod")

if ( not os.path.isfile ( modin_file ) ):
     print ( "Cannot find calibrated visibility file %s" % modin_file )
     exit  ( 1 )

if ( modout_file == modin_file ):
     print ( "Wrong inpout file %s -- the file name is expected to end with uvs.fits" % modin_file )
     exit  ( 1 )

band = modin_file[-5:-4]

print  ( "band= ", band ) 

if ( band == "L" ):
     dist_from_main = 90.0
     dist_from_comp = 15.0
     size           =  8.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "S" ):
     dist_from_main = 20.0
     dist_from_comp =  6.0
     size           =  3.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "C" ):
     dist_from_main = 12.0
     dist_from_comp =  5.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "F" ):
     dist_from_main = 10.0
     dist_from_comp =  4.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "X" ):
     dist_from_main = 10.0
     dist_from_comp =  4.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "U" ):
     dist_from_main = 10.0
     dist_from_comp =  5.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "K" ):
     dist_from_main = 10.0
     dist_from_comp =  4.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0

elif ( band == "Q" ):
     dist_from_main = 10.0
     dist_from_comp = 14.0
     size           =  2.0
     ratio          =  1.0
     angle          =  0.0


buf = read_file ( modin_file )

out     = []
out_com = []
out_com.append ( "! Input model generated by gen_input_model.py" )
out_com.append ( "! Input file: %s " % modin_file )

tot_flux_main = 0.0
tot_flux_comp = 0.0

nc_tot  = 0
nc_main = 0
nc_comp = 0
nc_est  = 0

for line in buf:
#    print ( "line= ", line )
    if ( "! Freq (Hz)     SpecIndex" in line ): continue
    if ( "!" in line ):
         out_com.append ( line )
         continue
    flux = float( line.split()[0] )
    nc_tot = nc_tot +  1
    if ( flux_main == 0.0 ):
         flux_main = flux
    if ( flux < flux_main*dyn_range ): continue
    dist = float( line.split()[1] )
    ang  = float( line.split()[2] )

    if ( abs(dist - dist_apr) < dist_from_comp ):
         if ( flux > flux_comp ):
              flux_comp = flux
              dist_comp = dist
              ang_comp  = ang 
              tot_flux_comp = tot_flux_comp + flux
         nc_comp = nc_comp + 1
         continue
    elif ( dist < dist_from_main ):
         tot_flux_main = tot_flux_main + flux
         nc_main = nc_main + 1

    tot_flux = tot_flux + flux

    if ( dist < dist_from_main or dist < dist_from_main  or "!" in line ):
         out.append ( "%9.6f  %9.4f  %9.2f" % ( flux, dist, ang ) )


out.sort ( key=lambda field: float(field.split()[0]), reverse=True )

for i in range(0,len(out)):
    flux_str = out[i].split()[0]
    flux     = float ( out[i].split()[0] )
    if ( flux >= est_range*flux_main ):
         nc_est = nc_est + 1
         out[i] = out[i].replace ( flux_str + " ", flux_str  + "v" )

out = out_com + out

print ( "Number of Clean comonents -- total: %4d, main: %3d, comp: %3d, est: %3d" % \
         ( nc_tot, nc_main, nc_comp, nc_est ) )
print ( "Total flux: %9.4f  main: %9.4f  comp: %9.4f" % ( tot_flux, tot_flux_main, tot_flux_comp ) )
print ( "rmod %s  ; sele I ; modelfit 12 ;  mapcol grey; mapl map,true" % modout_file )

out.append ( "!" )
out.append ( "! Total flux: %9.4f  main: %9.4f  comp: %9.4f" % ( tot_flux, tot_flux_main, tot_flux_comp ) )
out.append ( "!" )
out.append ( "! aprori estimation model: " )
out.append ( "!" )
out.append ( "%9.6fv  %9.4fv  %9.2fv  %8.2fv  %8.2fv  %8.2fv   %d" % 
              ( tot_flux_comp, dist_comp, ang_comp, size, ratio, angle, 1 ) )
out.append ( "!" )
out.append ( "! restore_dimap %s" % uvs_file )
out.append ( "! rmod %s  ; sele I ; modelfit 12 ;  mapcol grey; mapl map,true" % modout_file )
out.append ( "!" )
out.append ( "! /astrogeo/4C+35.28/src/model_fit.csh %s %s 10 > %s " % \
              ( uvs_file[:-9], modin_file, modin_file.replace("_fit.mod","_fit.txt") ) )

(ret,err) = write_file ( out, modout_file )
check_err_exe ( ret, err, "write_file" )

print ( "Wrote file %s" % modout_file )