#!/bin/env python

#import from system
import sys
import os
import math
import argparse
import time
import numpy as np

#safe pymetranet import
import import_pymetranet
from pymetranet import(PolarSweepSerializer, PolarSweep, PolarSweepInfo,
    PolarPpiData, ProductData, ProductDataPolar, ProductDataRect,
    ProductFile, ProductTable, MapSizeRect)

def calc_size(in_data: PolarPpiData, sweep: PolarSweep) -> MapSizeRect:
    x_y_size: int = in_data.num_gates * 2
    #x and y size is calculated using the number of gates but resolution is forcibly set to 0.25
    #for test/debug purposes
    return MapSizeRect(x_y_size, x_y_size, 0.25, 0.25)
    
def create_header(prod: ProductFile, in_data: PolarPpiData, sweep: PolarSweep, sweep_file_name: str) -> None:
    sweep_pid = os.path.basename(sweep_file_name)[0:0+3]
    params_pid = "PZ" + sweep_pid[2:]

    moment_name = "Z"
    base_file_name = os.path.basename(sweep_file_name)
    file_name_no_ext, file_ext = os.path.splitext(sweep_file_name)
    sweep_rank: int = int(file_ext[1:])
    idx_mid_ray: int = int(len(sweep.rays) / 2)
    elevation = sweep.rays[idx_mid_ray].get_startel_deg() #get start el of mid ray
    prd_title = "PPI_%s_%d@Elevation_%3.2f" % (moment_name, sweep_rank, elevation)
    param_time = base_file_name[3:3+9]
    param_pri = base_file_name[12:12+1]
    param_comp = base_file_name[13:13+1]
    param_format = "RECT"

    sweep_info: PolarSweepInfo = PolarSweepInfo(sweep)
    nyquist = sweep_info.get_velocity_nyquist()

    #automatically calculate size depending on gate width and range
    size: MapSizeRect = calc_size(in_data, sweep)
    print("calculated size:", size)

    param_comments = "product generated by test_mch_writeproduct_32bit.py pymetranet example"

    prod.add_header_info("pid", params_pid)
    prod.add_header_info("product", prd_title)
    prod.add_header_info("time", "%s%s%s" % (param_time, param_pri, param_comp))
    prod.add_header_info("format", param_format)
    prod.add_header_info("moment", "Z")
    prod.add_header_info("data_unit", "dBZ")
    prod.add_header_info("radar", sweep.sweepheader.radarname)
    prod.add_header_info("row", size.y_size)
    prod.add_header_info("column", size.x_size)
    prod.add_header_info("data_elevation", "%f" % elevation)
    prod.add_header_info("data_width", "%f" % sweep.sweepheader.gatewidth)
    prod.add_header_info("data_time", int(sweep.rays[0].rayheader.datetime / 100))
    prod.add_header_info("data_bits", "32")
    prod.add_header_info("total_sweep", sweep.sweepheader.totalsweep)
    prod.add_header_info("current_sweep", sweep.sweepheader.currentsweep)
    prod.add_header_info("nyquist", "%f" % nyquist)
    #compressed_bytes and uncompressed_bytes MUST be present, but it's not
    #necessary they are set, save_file method will take care of their value
    prod.add_header_info("compressed_bytes", "")
    prod.add_header_info("uncompressed_bytes", "")
    if param_format == "RECT":
        prod.add_header_info("rect_xres", "%f" % size.x_res)
        prod.add_header_info("rect_yres", "%f" % size.y_res)
        prod.add_header_info("data_flags", sweep.rays[0].rayheader.dataflags)
        prod.add_header_info("rad_lat", "%f" % sweep.sweepheader.radarlat)
        prod.add_header_info("rad_lon", "%f" % sweep.sweepheader.radarlon)
        prod.add_header_info("rad_height", "%f" % sweep.sweepheader.radarheight)
    if param_comments: #not empty
        prod.add_header_info("comment_size", len(param_comments))
        prod.add_header_info("comments",  param_comments)
        
    #example of usage of parameters, 'param_num' must be present, then
    #we must use 'param_name'/'param_value' for the first param,
    #'param_name2'/'param_value2' for the second param,
    #'param_name3'/'param_value3' for the third param, and so on
    prod.add_header_info("param_num", "0")
    #prod.add_header_info("param_name", "ZR_a")
    #prod.add_header_info("param_value", "443.5")
    #prod.add_header_info("param_name2", "ZR_b")
    #prod.add_header_info("param_value2", "1.2987")
    #prod.add_header_info("param_name3", "my_custom_param")
    #prod.add_header_info("param_value3", "123456")

    #//example of usage of tables, 'table_num' must be present, then
    #//we must use 'table_name'/'table_size' for the first table,
    #//'table_name2'/'table_size2' for the second table,
    #//'table_name3'/'table_size3' for the third table, and so on
    prod.add_header_info("table_num", "0")
    #prod.add_header_info("table_name", "Metranet_8bit")
    #prod.add_header_info("table_size", "1024")

def create_float_32bit_data(in_data: PolarPpiData, vectorized: bool=True) -> np.ndarray:
    if vectorized:
        return create_float_32bit_data_vectorized(in_data)
    
    buff_in = in_data.data
    buff_out = np.full(buff_in.shape, np.nan, dtype=np.float32)
    indices = zip(*np.where(np.isfinite(buff_in)))
    for i, j in indices:
        buff_out[i,j] = buff_in[i,j]

    return buff_out

def create_float_32bit_data_vectorized(in_data: PolarPpiData) -> np.ndarray:
    buff_in = in_data.data
    buff_out = np.full(buff_in.shape, np.nan, dtype=np.float32)

    mask = np.isfinite(buff_in)
    buff_out[mask] = buff_in[mask]

    return buff_out

def create_tables(prod: ProductFile) -> None:
    #here the table_num table_name, table_size in the header info
    #must be already present
    #we are creating a float 32bit product so conversion table named Metranet_8bit
    #is not necessary
    return

def create_data(prod: ProductFile, in_data: PolarPpiData, sweep: PolarSweep) -> None:
    #create output data in polar format
    polar: ProductDataPolar = ProductDataPolar(in_data.num_rays, in_data.num_gates, 4)

    #create encoded data
    polar.data = create_float_32bit_data(in_data)

    #if the output format is RECT, convert data to rect format
    #else keep data as is
    format = prod.find_header_info_value("format")
    if format == "RECT":
        #create a rect data and convert polar to rect
        x: int = prod.find_header_info_value("column")
        y: int = prod.find_header_info_value("row")
        x_res: float = float(prod.find_header_info_value("rect_xres"))
        y_res: float = float(prod.find_header_info_value("rect_yres"))
        data_bytes: int = int(prod.find_header_info_value("data_bits")) // 8
        rect: ProductDataRect = ProductDataRect(x, y, x_res, y_res, data_bytes)
        rect.polar2rect(polar, sweep.sweepheader.gatewidth)

        #set encoded rect data to the product
        prod.data = rect
    else: #polar
        #set encoded polar data to the product
        prod.data = polar

def adjust_header(prod: ProductFile) -> None:
    #no need to adjust, just verify with assert
    row_in_header: int = int(prod.find_header_info_value("row"))
    assert row_in_header == prod.data.num_rows
    column_in_header: int = int(prod.find_header_info_value("column"))
    assert column_in_header == prod.data.num_cols

def save(prod: ProductFile, sweep_file_name: str) -> None:
    #get base name and dir of sweep file name,
    #go up one level, get pid of product and
    #build full path of product file to be saved on disk
    sweep_base_file_name = os.path.basename(sweep_file_name)
    sweep_file_timestamp = param_time = sweep_base_file_name[3:3+9]
    sweep_file_priority = sweep_base_file_name[12:12+1]
    sweep_file_path_no_ext, sweep_file_ext = os.path.splitext(sweep_file_name)
    data_base_dir = os.path.abspath(os.path.join(os.path.dirname(sweep_file_name), ".."))
    prod_pid = prod.find_header_info_value("pid")
    out_prod_dir = os.path.join(data_base_dir, prod_pid)
    if not os.path.exists(out_prod_dir):
        print("creating output directory '%s' ..." % out_prod_dir)
        os.mkdir(out_prod_dir)
    out_prod_file_name = "%s%s%sL%s" % (prod_pid, sweep_file_timestamp, sweep_file_priority, sweep_file_ext)
    out_prod_full_path = os.path.join(out_prod_dir, out_prod_file_name)

    print("writing output product to '%s' ..." % out_prod_full_path)
    prod.save_file(out_prod_full_path, compress=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('infile', nargs='*', type=str)
    
    args = parser.parse_args()
    
    for file_name in args.infile:
        print("loading file %s..." % file_name)
        loaded_sweep = PolarSweepSerializer.load(file_name)
        print("sweep file %s successfully loaded!" % file_name)
        
        #transform 'mom'
        in_data = PolarPpiData()
        in_data.transform(loaded_sweep, mom_name="Z")

        #create product
        prod: ProductFile = ProductFile()
        create_header(prod, in_data, loaded_sweep, file_name)
        create_tables(prod)
        create_data(prod, in_data, loaded_sweep)
        adjust_header(prod)

        #save product to disk
        save(prod, file_name)
        
if __name__ == '__main__':
    main()
