#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 21 11:01:52 2017

@author: Neo

Retrieve the estimates of positions of local stations and the formal
uncertainties of these estimates from .lst file which is generated by
the program getpar.




"""

import numpy as np
import sys
from get_pv import get_dat, get_sta_pv
from write_pv import write_sta_pv
from plot_data import plot_sta_pv


# ------------------------------  FUNCTIONS  ---------------------------
def wp_sta_pos(staname, stalist, epo, X, Y, Z, X_err, Y_err, Z_err,
               U, E, N, U_err, E_err, N_err):
    '''Write and plot geocentric and topocentric positions for each station.
    '''
    eposta = get_dat(staname, stalist, epo)

    # For geocentric position
    [Xsta, Ysta, Zsta,
     Xsta_err, Ysta_err, Zsta_err] = get_sta_pv(
        staname, stalist, X, Y, Z, X_err, Y_err, Z_err)
    write_sta_pv(staname, "XYZ", eposta,
                 Xsta, Ysta, Zsta,
                 Xsta_err, Ysta_err, Zsta_err)
    plot_sta_pv(staname, "XYZ", eposta,
                Xsta, Ysta, Zsta,
                Xsta_err, Ysta_err, Zsta_err, 'p')

    # For topocentric position
    [Usta, Esta, Nsta,
     U_errsta, E_errsta, N_errsta] = get_sta_pv(
        staname, stalist, U, E, N, U_err, E_err, N_err)
    write_sta_pv(staname, "UEN", eposta,
                 Usta, Esta, Nsta,
                 U_errsta, E_errsta, N_errsta)
    plot_sta_pv(staname, "UEN", eposta,
                Usta, Esta, Nsta,
                U_errsta, E_errsta, N_errsta, 'p')


def strl2flta(X):
    '''transform the list of string type into an array of float type.
    '''
    return np.array(X).astype(float)


def read_lst(datafile):
    '''Retrieve the result from .lso file.

    Parameters
    ----------
    datafile : string
        name of data file

    Returns
    ----------
    None.
    '''

    # empty list for store data
    dbname = []
    stalist = []
    epo = []
    # XYZ
    X = []
    X_err = []
    Y = []
    Y_err = []
    Z = []
    Z_err = []
    # UEN
    U = []
    U_err = []
    E = []
    E_err = []
    N = []
    N_err = []

    for line in open(datafile, 'r'):
        if line[0] != '#':
            if line[:7] == 'STA_LCX':
                stalist.append(line[10:18].rstrip())
                dbname.append(line[20:30])
                epo.append(line[46:56])
                X.append(line[61:76])
                X_err.append(line[80:90])
                Y.append(line[95:110])
                Y_err.append(line[114:124])
                Z.append(line[129:144])
                Z_err.append(line[148:159])
            else:
                U.append(line[61:76])
                U_err.append(line[80:90])
                E.append(line[95:110])
                E_err.append(line[114:124])
                N.append(line[129:144])
                N_err.append(line[148:159])

    # string -> float, list -> array
    dbname = np.array(dbname)
    stalist = np.array(stalist)
    epo = strl2flta(epo)
    X = strl2flta(X)
    X_err = strl2flta(X_err)
    Y = strl2flta(Y)
    Y_err = strl2flta(Y_err)
    Z = strl2flta(Z)
    Z_err = strl2flta(Z_err)
    U = strl2flta(U)
    U_err = strl2flta(U_err)
    E = strl2flta(E)
    E_err = strl2flta(E_err)
    N = strl2flta(N)
    N_err = strl2flta(N_err)

    staset = set(stalist)
    for staname in staset:
        wp_sta_pos(staname, stalist, epo,
                   X, Y, Z, X_err, Y_err, Z_err,
                   U, E, N, U_err, E_err, N_err)


# Retrieve estimates.
if len(sys.argv) == 1:
    datafile = 'result/test.lst'
else:
    datafile = sys.argv[1]
read_lst(datafile)
# ------------------------------ END -----------------------------------
