from pprint import pprint as pp
import pickle
import numpy as np
import argparse
import yaml
import time
import sys
import numpy
import matplotlib.pyplot as plt
import os
import matplotlib.animation as animation
from pathlib import Path


def parse_args(args):
    parsed_args = None
    parser = argparse.ArgumentParser()
    parser.add_argument('scan_data_filename', type=str)
    parser.add_argument('pixel_resolution', type=float)
    parser.add_argument('--canvas_range', nargs="+", type=int)
    parser.add_argument('--animation', action='store_true')

    # add canvas_size and pixel_resolution arguments
    try:
        parsed_args = parser.parse_args()
    except Exception as e:
        raise RuntimeError("Could not parse args due to invalid/missing argument: {0}".format(e))

    if parsed_args is None:
        raise RuntimeError("")

    return parsed_args


def scan_iterate(xx,yy,platform_pos,range_bins,canvas_size,scan_data):
        mm_x = (xx - platform_pos[0])
        mm_y = (yy - platform_pos[1])
        mm_z = np.full(canvas_size, platform_pos[2], dtype='float16') # create array before?
        o = np.sqrt(np.square(mm_x)+np.square(mm_y)+np.square(mm_z))
        v = np.interp(o, range_bins, scan_data)
        return v


def main(args): # refactor this file so that it can be animated or not -- make it more vesatile
    # fig = plt.figure()
    # a = time.time()
    parsed_args = parse_args(args)

    with open(parsed_args.scan_data_filename,'rb') as f:
        data = pickle.load(f)

    # platform_pos = data['platform_pos']
    # with open(parsed_args.mocap_filename,'r') as f:
        # platform_pos =

    range_bins = data['range_bins'] # np.reshape() # (1056,)
    scan_data = data['scan_data']                  # (1287,1056)

    canvas_range = parsed_args.canvas_range
    pixel_resolution = parsed_args.pixel_resolution
    canvas_size = (int((canvas_range[1]-canvas_range[0])/pixel_resolution),
                   int((canvas_range[3]-canvas_range[2])/pixel_resolution))

    total_distance = 12.2
    curated_scan_data = scan_data[400:980]
    distance_interval = total_distance/len(curated_scan_data)

    # curated_range = 980-400 # 580

    platform_pos = np.zeros((len(curated_scan_data),3))

    platform_pos[:,0] = np.arange(-total_distance/2, total_distance/2, distance_interval)
    platform_pos[:,1] = 0
    platform_pos[:,2] = 1

    # pp(platform_pos)

    curated_platform_pos = platform_pos#[400:980]

    # pp(curated_platform_pos)


    # pp(curated_scan_data.shape)

    height = np.arange(canvas_range[0],canvas_range[1],pixel_resolution) # constructor with 120 times 0.12 steps from -3 to 3
    width = np.arange(canvas_range[2],canvas_range[3],pixel_resolution)

    xx, yy = np.meshgrid(width,height) # check direction...?
    output = np.zeros(canvas_size,dtype='complex64')
    output_array = []

    for i in range(len(curated_scan_data)):
        # pass
        v = scan_iterate(xx,yy,curated_platform_pos[i],range_bins,canvas_size,curated_scan_data[i])
        output += v
        # feed it the pulse each time...
        # if parsed_args.animation == True:
        #     abs_val_img = np.abs(output)
        #     img = plt.imshow(abs_val_img,animated=True)
        #     output_array.append([img])
        #     pp(i)

    if parsed_args.animation == True:
        # ani = animation.ArtistAnimation(fig, output_array, interval=25, repeat_delay=300)
        # ani.save('{0}.mp4'.format(Path(parsed_args.scan_data_filename).stem))
        # plt.show()
        pass
    else:
        output = np.flip(output,0)
        abs_val_img = np.abs(output)
        pp(abs_val_img)
        plt.imshow(abs_val_img)
        # pp(time.time()-a)
        plt.show()
        pp(platform_pos)


def printing(args):
    parsed_args = parse_args(args)
    with open(parsed_args.scan_data_filename,'rb') as f:
        data = pickle.load(f)
        pp(data)


if __name__ == "__main__":
    # a = time.time()
    main(sys.argv[1:])
    # printing(sys.argv[1:])
    # pp(time.time()-a)