#!/usr/bin/env python

import os
import sys
import textwrap

import toml
import click
from tabulate import tabulate
from click import confirm, style

import numpy as np

from libpipe import futils
from pspipe import __version__
from pspipe import settings, database


t_file = click.Path(exists=True, dir_okay=False)
t_dir = click.Path(exists=True, file_okay=False)


@click.group()
@click.version_option(__version__)
def main():
    ''' Add/list/remove obs_ids from a pspipe database'''


def load_settings(settings_file, config_modifiers=None, **additional_modifiers):
    try:
        s = settings.Settings.load_with_defaults(settings_file)
        if config_modifiers:
            mod = settings.Settings.load_from_string(config_modifiers)
            s = s + mod
    except toml.TomlDecodeError as e:
        print(f'Error parsing configuration: {e}')
        sys.exit(1)
    return s + additional_modifiers


@main.command('init')
@click.argument('template_name', type=str)
@click.argument('data_dir', type=str)
@click.argument('rev_name', type=str)
@click.argument('config_modifier', nargs=-1, required=False)
def init(template_name, data_dir, rev_name, config_modifier):
    ''' Initialize new configuration

    \b
    TEMPLATE_NAME: Template name (nenufar, hba or a12)
    DATA_DIR: Data directory
    REV_NAME: Image revision name
    CONFIG_MODIFIER: Optional configuration modifiers, of the form key=value
    '''
    if os.listdir(os.getcwd()):
        print('Current working directory is not empty, configuration file can not be created.')
        return

    settings_file = settings.get_template_path(template_name)
    print(f'Initializing configuration from {settings_file}')

    rev_name = rev_name.replace('.toml', '')
    conf_file = rev_name + '.toml'

    s = load_settings(settings_file, config_modifier, **{'data_dir': futils.abspath(data_dir),
                                                         'image': {'name': rev_name}})
    s.duplicate(conf_file, copy_parset_files=True)

    print(f'New configuration saved to {conf_file}')


@main.command('clone')
@click.argument('settings_file', type=t_file)
@click.argument('data_dir', type=str)
@click.argument('config_modifier', nargs=-1, required=False)
def clone(settings_file, data_dir, config_modifier):
    ''' Clone configuration and Obs IDs from an other pspipe config file

        \b
        CONFIG_MODIFIERS: Optional configuration modifiers, of the form key=value
    '''
    print(f'Cloning configuration from {settings_file}')
    s_in = load_settings(settings_file, config_modifier)
    s_out = s_in.duplicate(os.path.basename(settings_file), copy_parset_files=True, data_dir=futils.abspath(data_dir))

    db_in = database.Database(s_in)
    db_out = database.Database(s_out)
    for obs_id in db_out.list_obs_ids():
        db_out.new_obs_id_from_other(obs_id, db_in.get_obs_id_meta(obs_id))
        print(f'Adding obs_sd {obs_id}')


@main.command('new_rev')
@click.argument('settings_file', type=t_file)
@click.argument('rev_name', type=str)
@click.argument('config_modifier', nargs=-1, required=False)
def new_rev(settings_file, rev_name, config_modifier):
    ''' Create a new image revision

        \b
        CONFIG_MODIFIERS: Optional configuration modifiers, of the form key=value
    '''
    rev_name = rev_name.replace('.toml', '')
    conf_file = rev_name + '.toml'
    s = load_settings(settings_file, config_modifier)
    root_settings_file = s.get_root_settings_file()
    s.duplicate(conf_file, strip_parents_keys=True, default_settings=root_settings_file, image={'name': rev_name})
    print(f'New configuration saved to {conf_file}')


@main.command('list_revs')
@click.argument('settings_file', type=t_file)
@click.option('--obs_id', '-o', type=str, help='Obs ID')
@click.option('--img_rev', '-i', type=str, help='Image revision')
def list_revs(settings_file, obs_id, img_rev):
    ''' List all image and visibilities revisions
    '''
    s = load_settings(settings_file)
    db = database.Database(s)

    print(f"Image revisions in {s.data_dir}:")
    w_i = textwrap.TextWrapper(width=100, initial_indent='Obs: ', break_long_words=False)
    w_v = textwrap.TextWrapper(width=100, initial_indent='   Obs: ', subsequent_indent='   ',
                               break_long_words=False)

    for db_img_rev in db.list_image_revisions():
        if img_rev is not None and db_img_rev.name != img_rev:
            continue
        if obs_id and not db_img_rev.has_obs_id(obs_id):
            continue

        print(f'\n{style(db_img_rev.name, bold=True)}: {db_img_rev.get_path()}')
        print(w_i.fill(" ".join(db_img_rev.list_obs_ids())))
        for db_vis_rev in db_img_rev.list_vis_revisions():
            if obs_id is not None and not db_vis_rev.has_obs_id(obs_id):
                continue

            print(f' * {style(os.path.basename(db_vis_rev.get_path()), bold=True)}')
            print(w_v.fill(" ".join(db_vis_rev.list_obs_ids())))


def print_obs_ids(db, obs_ids):
    obs_metas = list(map(db.get_obs_id_meta, obs_ids))
    if not obs_metas:
        print('No obs_ids in database')
        return

    data = [obs_ids]

    for k in database.obs_id_info_keys:
        data.append([meta[k] for meta in obs_metas])

    data = list(map(list, zip(*data)))
    header = ['Obs_id'] + database.obs_id_info_keys

    print(tabulate(data, header))


@main.command('list_obs')
@click.argument('settings_file', type=t_file)
def list_obs(settings_file):
    ''' List obs_ids for the given image_rev/vis_rev
    '''
    s = load_settings(settings_file)
    vis_rev = database.VisRevision(s)
    if not os.path.exists(vis_rev.get_path()):
        print(f'Error: visibilities revision does not exist.')
        return

    print_obs_ids(database.Database(s), vis_rev.list_obs_ids())


@main.command('list_all_obs')
@click.argument('settings_file', type=t_file)
def list_all_obs(settings_file):
    ''' List obs_ids in the pspipe database
    '''
    db = database.Database(load_settings(settings_file))
    print_obs_ids(db, db.list_obs_ids())


def contains_duplicates(X):
    return len(np.unique(X)) != len(X)


def add(db, obs_id, ms_file_list, **kargs):
    if db.has_obs_id(obs_id):
        print(f'Error: obs_id {obs_id} already exist.')
        return

    if ms_file_list:
        ms_files = [k.strip().split() for k in open(ms_file_list).readlines() if k.strip()]
        flat_ms_files = [os.path.basename(ms_file) for sublist in ms_files for ms_file in sublist]

        if contains_duplicates(flat_ms_files):
            print('Error: All MS should have different names')
            return

        db.new_obs_id(obs_id, ms_files)
    else:
        db.new_obs_id(obs_id, **kargs)

    print(f'{obs_id} successfully created.')


@main.command('add_obs')
@click.argument('settings_file', type=t_file)
@click.argument('obs_id', type=str)
@click.option('--ms_file_list', '-m', type=t_file, help='MS file list')
@click.option('--int_time', '-i', type=float, help='Integration time (in second)', default=0)
@click.option('--total_time', '-t', type=float, help='Total time (in second)', default=0)
@click.option('--chan_width', '-w', type=float, help='Channel width in Hz', default=0)
def add_obs(settings_file, obs_id, ms_file_list, int_time, total_time, chan_width):
    ''' Add obs_ids to the given pspipe database
    '''
    db = database.Database(load_settings(settings_file))
    add(db, obs_id, ms_file_list, int_time=int_time, total_time=total_time, chan_width=chan_width)


@main.command('add_all_obs')
@click.argument('settings_file', type=t_file)
@click.argument('ms_file_lists', nargs=-1, type=t_file)
def add_all_obs(settings_file, ms_file_lists):
    ''' Add obs_ids from MSs list files
    '''
    db = database.Database(load_settings(settings_file))
    for ms_list in ms_file_lists:
        obs_id = os.path.basename(ms_list)
        add(db, obs_id, ms_list)


@main.command('set_obs_info')
@click.argument('settings_file', type=t_file)
@click.argument('obs_ids', type=str)
@click.option('--ms_file_list', '-m', type=str, help='MS file list')
@click.option('--int_time', '-i', type=float, help='Integration time')
@click.option('--total_time', '-t', type=float, help='Total time')
@click.option('--chan_width', '-w', type=float, help='Channel width in Hz',)
def set_obs_info(settings_file, obs_ids, ms_file_list, int_time, total_time, chan_width):
    ''' Set meta data info

    \b
    OBS_IDS: Obs ID(s), comma separated with bash-type wildcards support
    '''
    db = database.Database(load_settings(settings_file))
    obs_id_set = db.filter_obs_ids(obs_ids)

    if not obs_id_set:
        print(f'Error: No obs_ids matching {obs_ids} in database.')
        return

    for obs_id in obs_id_set:
        meta = db.get_obs_id_meta(obs_id)

        if ms_file_list:
            assert os.path.exists(ms_file_list)

            ms_files = [k.strip().split() for k in open(ms_file_list).readlines()]
            meta.set_ms_files(ms_files)
        else:
            d = {'int_time': int_time, 'total_time': total_time, 'chan_width': chan_width}
            for k, v in d.copy().items():
                if v is None:
                    del d[k]
            meta.update(d)

        meta.save()

        print(f'{obs_id} successfully modified.')


@main.command('add_rev_img')
@click.argument('settings_file', type=t_file)
@click.argument('obs_id', type=str)
@click.argument('fits_images', nargs=-1, type=t_file)
@click.option('--even', '-e', help='Fits images are even set', is_flag=True)
@click.option('--odd', '-o', help='Fits images are odd set', is_flag=True)
def add_rev_img(settings_file, obs_id, fits_images, even, odd):
    ''' Add FITS_IMAGES to the given pspipe database. An OBS_ID is required, which can be created with the 'add_obs' task.
    '''
    db = database.Database(load_settings(settings_file))
    if not db.has_obs_id(obs_id):
        print(f'Error: obs_id {obs_id} does not exist. Create with psdb add_obs {obs_id}')
        return

    if db.has_image_revisions(db.settings.image.name):
        img_rev = db.get_image_revision()
        if img_rev is None:
            img_rev = db.new_image_revision(db.settings.image.name)
    else:
        img_rev = db.new_image_revision(db.settings.image.name)

    if not img_rev.has_obs_id(obs_id):
        img_meta = img_rev.new_obs_id(obs_id)
    else:
        img_meta = img_rev.get_meta(obs_id)

    files = list(map(futils.abspath, fits_images))
    img_meta.add_fits(files, even=even, odd=odd, verbose=True)
    img_meta.save()


@main.command('list_rev_img')
@click.argument('settings_file', type=t_file)
@click.option('--obs_ids', '-o', type=str, help='Only for given obs ID(s)')
def list_rev_img(settings_file, obs_ids):
    ''' Retrieve FITS images statistics
    '''
    db = database.Database(load_settings(settings_file))
    if not db.has_image_revisions(db.settings.image.name):
        print(f'Error: image revision does not exist.')
        return

    if obs_ids is not None:
        obs_id_set = db.filter_obs_ids(obs_ids)
        if not obs_id_set:
            print(f'Error: No obs_ids matching {obs_ids} in database.')
            return

    img_rev = db.get_image_revision()
    stats = {}

    for obs_id in img_rev.list_obs_ids():
        if obs_ids is not None and obs_id not in obs_id_set:
            continue
        meta = img_rev.get_meta(obs_id)
        if meta is None:
            continue
        stats[obs_id] = meta.get_fits_stat()

    if len(stats) == 0:
        return

    all_types = sorted(set(t for k in stats.values() for t in k.keys()))

    data = []
    for obs_id, stat in stats.items():
        data.append([obs_id] + [stat.get(k, 0) for k in all_types])

    header = ['Obs_id'] + all_types

    print(tabulate(data, header))


@main.command('remove_rev')
@click.argument('settings_file', type=t_file)
@click.option('--obs_ids', '-o', type=str, help='Only for given obs ID(s)')
def remove_rev(settings_file, obs_ids):
    ''' Remove an image revision from the given pspipe database
    '''
    db = database.Database(load_settings(settings_file))

    if obs_ids is not None:
        obs_id_set = db.filter_obs_ids(obs_ids)
        if not obs_id_set:
            print(f'Error: No obs_ids matching {obs_ids} in database.')
            return

    img_rev = db.get_image_revision()

    if img_rev is None:
        print(f'Error: image revision {db.settings.image.name} does not exist.')
        return

    if obs_ids is not None:
        for obs_id in img_rev.list_obs_ids():
            if obs_id in obs_id_set and confirm(style(f'Removing ObsID {obs_id} in revision {img_rev.name} ?',
                                                      fg='yellow')):
                img_rev.remove(obs_id=obs_id)
    else:
        if confirm(style(f'Removing {img_rev.name} ?', fg='yellow')):
            img_rev.remove()
            print(f'Image revision {db.settings.image.name} successfully removed.')
        else:
            print('No changes made.')


@main.command('remove_obs')
@click.argument('settings_file', type=t_file)
@click.argument('obs_ids', type=str)
def remove_obs(settings_file, obs_ids):
    ''' Remove an Obs ID from the given pspipe database
    '''
    db = database.Database(load_settings(settings_file))
    obs_id_set = db.filter_obs_ids(obs_ids)

    if not obs_id_set:
        print(f'Error: No obs_ids matching {obs_ids} in database.')
        return

    if confirm(style(f'Removing obs_ids: {",".join(obs_id_set)} ?', fg='yellow')):
        for obs_id in obs_id_set:
            db.remove_obs_id(obs_id)
            print(f'{obs_id} successfully removed.')
    else:
        print('No changes made.')


if __name__ == '__main__':
    main()
