from ObservationSim import Observation
from _util import parse_args
from datetime import datetime
import os
import numpy as np
import galsim
import yaml

import gc
gc.enable()

def Pointing(config, pointing_filename, data_dir):
    pointing_file = os.path.join(data_dir, pointing_filename)
    f = open(pointing_file, 'r')
    for _ in range(1):
        header = f.readline()
    iline = 0
    pRA = []
    pDEC = []
    for line in f:
        line = line.strip()
        columns = line.split()
        pRA.append(float(columns[0]))
        pDEC.append(float(columns[1]))
    f.close()
    pRA = np.array(pRA)
    pDEC = np.array(pDEC)

    # Create calibration pointings
    # NOTE: temporary implementation
    ncal = config['obs_setting']['np_cal']
    pointing_type = ['MS']*len(pRA)
    pRA = np.append([pRA[0]]*ncal, pRA)
    pDEC = np.append([pDEC[0]]*ncal, pDEC)
    pointing_type = ['CAL']*ncal + pointing_type

    # Calculate starting time(s)
    # NOTE: temporary implementation
    t0 = datetime(2021, 5, 25, 12, 0, 0)
    t = datetime.timestamp(t0)
    timestamp_obs = []
    delta_t = 10 # Time elapsed between exposures (minutes)
    for i in range(len(pointing_type)):
        timestamp_obs.append(t)
        if pointing_type[i] == 'CAL':
            t += 3 * delta_t * 60 # 3 calibration exposure
        elif pointing_type[i] == 'MS':
            t += delta_t * 60
    timestamp_obs = np.array(timestamp_obs)
    pointing_type = np.array(pointing_type)
    return pRA, pDEC, timestamp_obs, pointing_type

def MakeDirectories(work_dir, run_name, nPointings, pRange=None):
    if not os.path.exists(work_dir):
        try:
            os.makedirs(work_dir, exist_ok=True)
        except OSError:
            pass
    imgDir = os.path.join(work_dir, run_name)
    if not os.path.exists(imgDir):
        try:
            os.makedirs(imgDir, exist_ok=True)
        except OSError:
            pass
    prefix = "MSC_"
    for pointing_ID in range(nPointings):
        if pRange is not None:
            if pointing_ID not in pRange:
                continue
        fname=prefix + str(pointing_ID).rjust(7, '0')
        subImgDir = os.path.join(imgDir, fname)
        if not os.path.exists(subImgDir):
            try:
                os.makedirs(subImgDir, exist_ok=True)
            except OSError:
                pass

def runSim():
    """
    Main simulation call.
    """
    args = parse_args()
    if args.config_dir is None:
        args.config_dir = ''
    args.config_dir = os.path.abspath(args.config_dir)
    args.config_file = os.path.join(args.config_dir, args.config_file)
    with open(args.config_file, "r") as stream:
        try:
            config = yaml.safe_load(stream)
            for key, value in config.items():
                print (key + " : " + str(value))
        except yaml.YAMLError as exc:
            print(exc)
    config["obs_setting"]["image_rot"] = float(config["obs_setting"]["image_rot"])*galsim.degrees

    if args.data_dir is not None:
        config['data_dir'] = args.data_dir
    if args.work_dir is not None:
        config['work_dir'] = args.work_dir

    pRA, pDEC, timestamp_obs, pointing_type = Pointing(config=config, pointing_filename=config['pointing_file'], data_dir=config['data_dir'])

    MakeDirectories(work_dir=config['work_dir'], run_name=config['run_name'], nPointings=len(pRA), pRange=config['obs_setting']['run_pointings'])
    if "run_chips" in config["obs_setting"]:
        run_chips = config["obs_setting"]["run_chips"]
    else:
        run_chips = None

    # from Config import ConfigDir
    # path_dict = ConfigDir(config=config, work_dir=config["work_dir"], data_dir=config["data_dir"])
    # for key, value in path_dict.items():
    #     print (key + " : " + str(value))

    obs = Observation(config=config, work_dir=config['work_dir'], data_dir=config['data_dir'])
    if config["pointing_file"] is None:
        obs.runExposure(chips=run_chips)
    else:
        obs.runExposure_MPI_PointingList(
            ra_cen=pRA, 
            dec_cen=pDEC, 
            pRange=config["obs_setting"]["run_pointings"], 
            timestamp_obs=timestamp_obs,
            pointing_type=pointing_type,
            exptime=config["obs_setting"]["exp_time"],
            use_mpi=config["run_option"]["use_mpi"],
            chips=run_chips
            )
    print("run finished")

if __name__=='__main__':
    runSim()

#############################################
# Testing run one exposure (NOT using MPI)
# ipoint = 2
# pointRa = float('{:8.4f}'.format(pRA[ipoint]))
# pointDec = float('{:8.4f}'.format(pDEC[ipoint]))

# obs = Observation(work_dir=work_dir, data_dir=data_dir)
# obs.runExposure(ra_cen=pointRa, dec_cen=pointDec, pointing_ID=ipoint, oneChip=16)

#############################################
# Testing run pointing list (using MPI)
# obs = Observation(work_dir=work_dir, data_dir=data_dir)
# # obs.runExposure_MPI_PointingList(ra_cen=pRA, dec_cen=pDEC, pRange=pRange, timestamp_obs=timestamp_obs, pointing_type=pointing_type)
# obs.runExposure_MPI_PointingList(ra_cen=pRA, dec_cen=pDEC, pRange=pRange, timestamp_obs=timestamp_obs, pointing_type=pointing_type, exptime=exp_time)
