#!/usr/bin/env python3

import pandas as pd
import numpy as np
import os
import argparse
import sys
import subprocess
import csv
import matplotlib.pyplot as plt

# For some gneric utiltities: json_read, json_write, timer
import montera.gutil as gutil
import montera.pos_lib as pos_lib
pos_lib.ant2ap_mode = 'single_file'


def main():
    total_time = gutil.timer('total')

    parser = argparse.ArgumentParser(
        description='Analyze swapped positions',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        #epilog="""Under the hood, we run the apps/pos_in_one.cpp code to create
        #.csv files, which we then read. Data is ploted to .png files."""
    )
    parser.add_argument(
        '-F', "--folders",
        metavar="FOLDERS",
        nargs='+',
        default=[
            #'240216_first_pos_not_used',
            '240217_first_pos',
            '240218_first_pos',
            '240219_swap_pos',
            '240220_swap_pos',
            #'240222_first_pos_N_not_used',
            '240223_first_pos_N',
            '240224_first_pos_N',
            #'240225_first_pos_N_include_shift',
            '240225_first_pos_N_after_shift',
            #'240226_first_pos_N_not_used',
            '240305_swap_pos_N',
            '240306_swap_pos_N',
            '240307_swap_pos_N',
            ],
        help="Base folder names. "
             "The first in the list will be used as the reference."
    )
    parser.add_argument(
        "--truth_csv",
        default='center.csv',
        help="CSV to use for Truth",
    )
    parser.add_argument(
        "--dut_dir",
        default = "BLDG2BD",
        help="directory with DUT T04 data. It will be used as "
             "FOLDER[n]/data/gnss/DUT_DIR[n]/. If only a single DUT_DIR is "
             "provided, then it will be used for all FOLDERS",
    )
    outp = parser.add_argument_group('output folder options')
    outp.add_argument(
        "--data_out_dir",
        default="gen_data",
        help="Base filename for file I/O. Inputs are .T04. "
        "Outputs are .png, .csv"
    )
    outp.add_argument(
        "--log_dir",
        default="logs",
        help="where to put the log (.txt) from running snr_diff. "
             "It will be used as FOLDER/LOG_DIR/",
    )
    outp.add_argument(
        "--fig_dir",
        default="figures",
        help="where to put the PNG. "
             "It will be used as FOLDER/FIG_DIR/",
    )
    parser.add_argument(
        "--force",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Load data from a CSV and json files rather than t04.",
    )
    parser.add_argument(
        "--plot",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Create plots"
        )

    args = parser.parse_args()

    reader = csv.DictReader(open(args.truth_csv))
    center = next(reader) 
    c_lat = float(center['lat'])
    c_lon = float(center['lon'])
    c_hgt = float(center['hgt'])
    
    os.makedirs(args.log_dir, exist_ok=True)
    # ------------ File name mangling
    records = [{
            'path': 'center',
            'antenna': 2,
            'lat': c_lat,
            'lon': c_lon,
            'hgt': c_hgt,
            'easting': 0,
            'northing': 0,
            'dh': 0,
            'distance': 0,
            }]
    for n, folder in enumerate(args.folders):
        print(f"============ {folder} {args.dut_dir} =====================")
        #data_dir = os.path.join(
        #        folder, args.data_out_dir, 'position', dut_dir)
        #print(f"{data_dir=}")
        log_fname = os.path.join(args.log_dir, folder + 'log.txt')

        fname = os.path.join(folder, 'gen_data', 'position', args.dut_dir,
                                  'BLDG2BD_pos_weighted_stats.csv')
        print(f'{fname=}')
        have_data = os.path.isfile(fname)
        if args.force or not have_data:
            cmd_list = ['pos_in_one_stats',
                        folder,
                        args.dut_dir,
                        '--force',
                        '--ne_tight_lim',
                        ]
            #print("---------------------------------------------")
            print(' '.join(cmd_list))

            f = open(log_fname, 'w')
            p = subprocess.run(
                cmd_list, stderr=subprocess.STDOUT, stdout=f
            )
            f.close()
            if p.returncode != 0:
                print(f'Warning pos_in_one_stats return code {p.returncode}')

        df = pd.read_csv(fname)
        for index, row in df.iterrows():
            d_lat = row['w_mean_lat']
            d_lon = row['w_mean_lon']
            d_hgt = row['w_mean_hgt']
            dist = pos_lib.dist_m(c_lat, c_lon, d_lat, d_lon)
            enu =  pos_lib.llh2enu(
                                        np.array([c_lat, c_lon, c_hgt]),
                                        np.array([d_lat, d_lon, d_hgt]),
                                        is_rad=False,
                                        is_ref_rad=False,
                                        )
            easting = enu[0][0]
            northing = enu[1][0]
            dh=enu[2][0]
            record = {
                    'path': folder,
                    'antenna': float(row.antenna),
                    'lat': d_lat,
                    'lon': d_lon,
                    'hgt': d_hgt,
                    'easting': easting,
                    'northing': northing,
                    'dh': dh,
                    'distance': dist,
                    }
            records.append(record)


    df = pd.DataFrame(records)
    print(df)
    print(df.info())

    # move folder and dut to first two columns
    print("Saving pos_swap.csv")
    df.to_csv('pos_swap.csv', index=False)

    if args.plot:
        #df['c']=(df.antenna+1.0)/3.0
        #print(df)

        m = np.max(np.array([df.easting.abs().max(), df.northing.abs().max()]))

        fig, ax = plt.subplots(1, 1, figsize=[12,5])
        """
        plt.scatter(df.easting, df.northing,
                    s=40,
                    #c=(df.antenna +0.5) / 3.5,
                    c=(df.antenna +0.1) / 2,
                    marker='+',
                    linewidth=0.5,
                    edgecolor=None,
                    )
        plt.colorbar()
        """
        df.query("antenna == 0").plot('easting', 'northing',
                                      ax=ax,
                                      marker='+', color='b',
                                      linewidth=0,
                                      label='Ant 0')
        df.query("antenna == 1").plot('easting', 'northing',
                                      ax=ax,
                                      marker='x', color='r',
                                      linewidth=0,
                                      label='Ant 1')
        #plt.legend()
        plt.xlim([-m, m])
        plt.ylim([-0.1, 0.1])
        plt.xlabel('Easting (m)')
        plt.ylabel('Northing (m)')
        plt.title('CSNMR-2963_BD992_Z3R_Swap_Positions')
        fig_fname = 'pos_swap.png'
        print(f"saving {fig_fname}")
        plt.savefig(fig_fname)
        plt.close(fig)

    total_time.dt_print()


if __name__ == '__main__':
    main()