#!/usr/bin/python
# Simulate the known radio sources during an entire day's observing
# Calculate visibilities and compare them to the observed visibilities
#
# Copyright (c) Tim Molteno 2013-2019. tim@elec.ac.nz
#
import traceback
import argparse
import math
import utc

import numpy as np
from multiprocessing import Pool


import datetime

import matplotlib.pyplot as plt
from tart.util import angle

import forward_map

from .visibility import *
from . import location
from ..operation import observation


def process_file(
    dirname, fname, orientation_degrees, spacing, sample_duration, simulate_vis=True
):
    try:
        fpath = os.path.join(dirname, fname)
        nobs = observation.Observation_Load(fpath)
        date = nobs.timestamp.date()
        sec = (
            nobs.timestamp - datetime.datetime.combine(date, datetime.time.min)
        ).total_seconds()

        if simulate_vis == True:
            utc_date = nobs.timestamp.replace(tzinfo=utc.UTC)
            antenna_locations = antenna_locations = nobs.ant_positions

            sky, baselines, visibility = forward_map.forward_map(
                utc_date,
                location.Dunedin,
                horizon=5.0,
                antenna_locations=antenna_locations,
                sample_duration=sample_duration,
            )  # Remove things below an elevation of five degrees
        else:
            sky = False
            baselines = False
            visibility = [False]

        hours = sec / 3600.0
        v = Visibility(nobs)
        return date, hours, v, visibility[0]
    except Exception as inst:
        print(type(inst))  # the exception instance
        print(inst.args)  # arguments stored in .args
        print(inst)
        print(sample_duration)
        print(traceback.format_exc())
        return None, None, None, None
    except KeyboardInterrupt:
        print(inst)
        return None, None, None, None


def get_correlation(
    dirname, orientation_degrees, spacing, sample_duration, simulate_vis=True
):
    files = sorted(os.listdir(dirname))
    vis = []
    simvis = []
    times = []

    if True:
        p = Pool()

        resultList = []

        try:
            for fname in files:
                if "data.pkl" in fname:
                    resultList.append(
                        p.apply_async(
                            process_file,
                            (
                                dirname,
                                fname,
                                orientation_degrees,
                                spacing,
                                sample_duration,
                                simulate_vis,
                            ),
                        )
                    )

            p.close
            p.join

            for sv_thread in resultList:
                [date, hours, v, sv] = sv_thread.get()  # timeout=2000
                if date != None:
                    times.append(hours)
                    vis.append(v)
                    simvis.append(sv)
                    # print("%g %g %g" % (times[-1], vis[-1], simvis[-1]))
            p.terminate()
        except KeyboardInterrupt:
            print("control-c pressed")
            p.terminate()

    else:
        for fname in files:
            if "data.pkl" in fname:
                [date, hours, v, sv] = process_file(
                    dirname,
                    fname,
                    orientation_degrees,
                    spacing,
                    sample_duration,
                    simulate_vis,
                )

                times.append(hours)
                vis.append(v)
                simvis.append(sv)

                print("%g %s %s" % (times[-1], v.v, simvis))

    # vis = np.array(vis)
    simvis = np.array(simvis)
    return date, times, vis, simvis


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Simulate the Sky and calculate simulated visibilities."
    )
    parser.add_argument(
        "--data-dir", required=True, help="The filesystem path for the telescope data."
    )
    parser.add_argument(
        "--orientation-angle",
        required=True,
        type=float,
        default=-22.0,
        help="The orientation of the array (degrees).",
    )
    parser.add_argument(
        "--optimize",
        action="store_true",
        help="Optimize the correlation between the simulated and measured visibilities.",
    )
    parser.add_argument(
        "--spacing", type=float, default=0.5, help="The spacing between antennas."
    )
    parser.add_argument(
        "--sample-duration",
        type=float,
        default=1e-4,
        help="The length of time to simulate data for.",
    )
    args = parser.parse_args()
    dirname = args.data_dir
    sample_duration = args.sample_duration
    orientation_degrees = args.orientation_angle
    spacing = args.spacing

    if args.optimize:
        for theta in np.arange(
            orientation_degrees - 10.0, orientation_degrees + 10.0, 1.0
        ):
            for r in np.arange(spacing - 0.03, spacing + 0.03, 0.01):

                date, times, vis, simvis = get_correlation(
                    dirname, theta, r, sample_duration
                )
                corr = sum(vis * simvis)
                print("Correlation @%f %f %f" % (theta, r, corr))
    else:
        date, times, vis, simvis = get_correlation(
            dirname, orientation_degrees, spacing, sample_duration
        )
        corr = sum(vis * simvis)
        print("Correlation @%f %f %f" % (orientation_degrees, spacing, corr))

        plt.subplot(2, 1, 1)
        title = "Visibility Data from %s (th=%f, dx=%f)" % (
            str(date),
            orientation_degrees,
            spacing,
        )
        plt.title(title)
        plt.plot(times, vis, ".")
        plt.ylabel("Measured V")
        plt.xlabel("Hours since UTC ")
        plt.subplot(2, 1, 2)
        plt.plot(times, simvis, ".")
        plt.ylabel("Simulated V")
        plt.xlabel("Hours since UTC ")
        plt.savefig(
            os.path.join(
                dirname,
                "simulated_visibility_%f_%f.pdf" % (orientation_degrees, spacing),
            )
        )
        plt.show()
