#! /usr/bin/python3
"""
Apply median filter to post ATS data with coseismic offsets already corrected.
"""
import multiprocessing as mp
import numpy as np
import os
import sys
from datetime import datetime
from collections import deque
from itertools import repeat

USAGE_MSG = "Usage: python reconstruct.py [postATS CleanTrend Directory]"
# Folder containing grids and the WNAM_SOPAC_neu_timeSerie_Length_corrDynDat_YYYYMMDD.dat active station list
GRID_FOLDER = "displGrid/"

# Multiprocessing process limit (use half of all available or 1)
if 'sched_getaffinity' in dir(os):
    # cpu count available to current process (linux only)
    PROCESS_LIMIT = len(os.sched_getaffinity(0)) // 2 or 1
else:
    PROCESS_LIMIT = os.cpu_count() // 2 or 1


def neuC(source, series_type, station):
    """ReconstructCo.py's neuC refers to a directory with coseismic
     offsets reconstructed and a median filter applied. Create the equivalent folder
     using input data with offsets already corrected.
     Take a station name as input and look for its file in the source directory.

    Args:
        source (str): Top level directory for neu and neuC folders.
        series_type (str): Clean_Trend or Filter_Trend.
        station (str): Lowercase station name whose file will be read.
    """
    print("Applying median for " + station)
    timeseries_file = source + '/neu/' + station + series_type + '.neu'
    if os.path.isfile(timeseries_file):
        xs2 = np.genfromtxt(timeseries_file, usecols=(0, 3, 4, 5, 6, 7, 8))

        # apply median to n,e,u
        for i in range(1,4):
            xs2[:, i] = medianSlidingWindow(xs2[:, i])

        with open(source + '/neuC/' + station + series_type + 'mFCo.neu', 'w') as TSMedian:
            for row in xs2:
                TSMedian.write('%9.4lf %9.2lf %9.2lf %9.2lf %8.2lf %8.2lf %8.2lf \n' % (
                row[0], row[1], row[2], row[3], row[4], row[5], row[6]))

    else:
        print("Not a file: " + timeseries_file)


def get_station_list():
    for f in sorted(os.listdir(GRID_FOLDER), reverse=True):
        if f.startswith('Station_Locations_'):
            print('Active station list: ', f)
            return f


def medianSlidingWindow(arr):
    # hardcode window size to 3
    if len(arr) < 2: 
        return arr

    window = deque([0, arr[0], arr[1]], maxlen=3)
    res = [sorted(window)[1]]

    for i in range(2, len(arr)):
        window.append(arr[i])
        res.append(sorted(window)[1])

    window.append(0)
    res.append(sorted(window)[1])

    return res


if __name__ == "__main__":
    start = datetime.now()

    # Take exactly 1 arg to use as source folder
    if len(sys.argv) == 2:
        source = sys.argv[1]
    else:
        print(USAGE_MSG)
        sys.exit(1)

    if "Clean" in source:
        series_type = "CleanTrend"
    else:
        series_type = "FilterTrend"

    Stations = []
    if os.path.isfile(source + '/neu/.DS_Store'):
        os.remove(source + '/neu/.DS_Store')

    # Create neuC folder
    if not os.path.exists(source + '/neuC'):
        os.mkdir(source + '/neuC')

    # Active list of stations
    active_list = get_station_list()
    if not active_list:
        print("Couldn't find the active station list")
        sys.exit()
    with open(GRID_FOLDER + active_list, 'r') as S:
        for line in S:
            Stations.append(line.split()[0])

    # Call neuC() on each station in Stations
    if sys.version_info >= (3,3):
        # Run multi-core to save precious seconds
        with mp.Pool(processes=PROCESS_LIMIT) as pool:
            iterable = zip(repeat(source), repeat(series_type), Stations)
            pool.starmap(neuC, iterable)
    
    else:
        print("Python version < 3.3")
        for s in Stations:
            neuC(source, series_type, s)

    print("Duration: " + str(datetime.now() - start)[:-4])