#!/usr/bin/env python3
# Author:  Michael Heflin
# Date:  September 12, 2014
# Organization:  JPL, Caltech

prolog="""
**PROGRAM**
    staBreak.py
      
**PURPOSE**
    | Search for breaks using an Ftest
    | All possible break locations are tested
    | Ftest is applied to break which minimizes CHI^2/DOF
    | If Ftest is passed that break is kept
    | Search continues until no remaining break location passes Ftest
"""
epilog="""
**EXAMPLE**
    staBreak.py -i ALGO.series -o ALGO.break --ftest 200
               
**COPYRIGHT**
    | Copyright 2014, by the California Institute of Technology
    | United States Government Sponsorship acknowledged
    | All rights reserved

**AUTHORS**
    | Michael Heflin
    | Jet Propulsion Laboratory
    | California Institute of Technology
    | Pasadena, CA, USA

Keywords: time series, reference frame, gdcat, break search
"""

import os
import argparse
import numpy as np
import time
import calendar

def _getParser():
    parser = argparse.ArgumentParser(description=prolog,epilog=epilog,
                            formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-i', action='store', dest='series',required=True,help='output gdcov file')
    parser.add_argument('-o', action='store', dest='output',required=True,help='output gdcov file')
    parser.add_argument('--ftest', action='store',dest='ftest',help='ftest value')
    return parser

def main():

    # Read command line arguments
    parser = _getParser()
    results = parser.parse_args()

    # Initialize
    T = []
    D = [] 
    N = [] 
    E = [] 
    V = [] 
    I = []
    L = []
    S = []

    # Read time series
    inFile = open(results.series,'r')
    line = inFile.readline()
    while line:
        test = line.split()
        T.append(test[0])
        N.append(test[2])
        E.append(test[1])
        V.append(test[3])
        D.append(float(test[10]))
        line = inFile.readline()
    inFile.close()

    # Initialize parameters
    ndat = len(T)
    I.append(ndat-1)
    A = np.zeros((ndat,2))
    A[0:ndat,0] = np.ones(ndat)
    A[0:ndat,1] = np.array(T,dtype=float) - 2015.0
    B = np.zeros((ndat,3))
    B[0:ndat,0] = np.array(N,dtype=float)
    B[0:ndat,1] = np.array(E,dtype=float)
    B[0:ndat,2] = np.array(V,dtype=float)
    a, b, c, d = np.linalg.lstsq(A,B,rcond=-1)
    c0 = b[0]+b[1]+b[2]/4
    p0 = 2

    # Search for breaks
    search = 1
    p1 = p0 + 1
    while search:
        A = np.hstack((A,np.zeros((ndat,1))))
        fmax = 0
        for i in range(0,ndat):
            A[i,p1-1] = 1
            fit = 1
            for j in range(0,len(I)):
                if (i == I[j]):
                    fit = 0
            if (fit == 1):
                a, b, c, d = np.linalg.lstsq(A,B,rcond=-1)
                c1 = b[0]+b[1]+b[2]/4
                f = ((c0-c1)/c1)*((ndat-p1)/(p1-p0))
                if (f > fmax):
                    imax = i
                    fmax = f
                    cmax = c1
                    Z = np.copy(A)
        if (fmax > float(results.ftest)):
            p0 = p1
            p1 = p1 + 1
            c0 = cmax
            I.append(imax) 
            A = np.copy(Z)
            list = D[imax+1]
            L.append(list)
        else:
            search = 0

    # Sort breaks
    L.sort(key=int)
    for i in range(0,len(L)):
        list = time.strftime("%Y %m %d %H %M %S",time.gmtime(calendar.timegm(time.strptime("2000JAN01 12:00:00","%Y%b%d %H:%M:%S"))+L[i]))
        S.append(str(list))

    # Remove existing file
    if os.path.exists(results.output):
        os.remove(results.output)

    # Print breaks
    if (len(S) > 0):
        outFile = open(results.output,'w')
        site = results.series[0:results.series.find('.')]
        for i in range(0,len(S)):
            print("{:s} {:s}".format(site,S[i]),file=outFile)
        outFile.close()
 
if __name__ == '__main__':
    main()