# Copyright (C) 2018 Cancer Care Associates

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version (the "AGPL-3.0+").

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License and the additional terms for more
# details.

# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

# ADDITIONAL TERMS are also included as allowed by Section 7 of the GNU
# Affero General Public License. These additional terms are Sections 1, 5,
# 6, 7, 8, and 9 from the Apache License, Version 2.0 (the "Apache-2.0")
# where all references to the definition "License" are instead defined to
# mean the AGPL-3.0+.

# You should have received a copy of the Apache-2.0 along with this
# program. If not, see <http://www.apache.org/licenses/LICENSE-2.0>.


import numpy as np

from pymedphys_utilities.utilities import get_filepath, get_gantry_tolerance

from pymedphys_mudensity.mudensity import get_grid


from pymedphys_databases.msq import multi_fetch_and_verify_mosaiq
from pymedphys_databases.delivery import DeliveryDatabases

from .analyse import calc_comparison, plot_results


def get_mappings(index, file_hashes):
    patient_grouped_fields = dict()
    field_id_grouped_hashes = dict()

    for file_hash in file_hashes:
        delivery_details = index[file_hash]['delivery_details']
        patient_id = delivery_details['patient_id']
        field_id = delivery_details['field_id']

        if patient_id not in patient_grouped_fields:
            patient_grouped_fields[patient_id] = set()

        patient_grouped_fields[patient_id].add(field_id)

        if field_id not in field_id_grouped_hashes:
            field_id_grouped_hashes[field_id] = []

        field_id_grouped_hashes[field_id].append(file_hash)

    return patient_grouped_fields, field_id_grouped_hashes


def group_consecutive_logfiles(file_hashes, index):
    times = np.array([
        index[key]['local_time']
        for key in file_hashes
    ]).astype(np.datetime64)

    sort_reference = np.argsort(times)
    file_hashes = file_hashes[sort_reference]
    times = times[sort_reference]

    hours_4 = np.array(60 * 60 * 4).astype(np.timedelta64)
    split_locations = np.where(np.diff(times) >= hours_4)[0] + 1

    return np.split(file_hashes, split_locations)


def assert_array_agreement(unique_logfile_gantry_angles, mosaiq_gantry_angles,
                           allowed_deviation):
    difference_matrix = np.abs(
        unique_logfile_gantry_angles[:, None] - mosaiq_gantry_angles[None, :])
    agreement_matrix = difference_matrix <= allowed_deviation
    row_agreement = np.any(agreement_matrix, axis=1)
    at_least_one_agreement = np.all(row_agreement)

    assert at_least_one_agreement, (
        'There is a logfile gantry angle that deviates by more than {} degrees'
        ' from the Mosaiq control points. Unsure how to handle this.\n\n'
        'Logfile: {}\nMosaiq: {}\nDifference Matrix:\n{}\n'
        'Agreement Matrix:\n{}'.format(
            allowed_deviation, unique_logfile_gantry_angles,
            mosaiq_gantry_angles, difference_matrix, agreement_matrix))


def get_field_id_from_logfile_group(index, logfile_group):
    field_ids = []

    for logfile_hash in logfile_group:
        field_ids.append(index[logfile_hash]['delivery_details']['field_id'])

    assert len(np.unique(field_ids)) == 1

    field_id = field_ids[0]

    return field_id


def calc_normalisation(mosaiq_delivery_data):
    mosaiq_delivery_data.m
    all_gantry_angles = mosaiq_delivery_data.mudensity
    mosaiq_gantry_angles = np.unique(mosaiq_delivery_data.gantry)
    number_of_gantry_angles = len(mosaiq_gantry_angles)

    normalisation = np.sum(all_gantry_angles) / number_of_gantry_angles

    return normalisation


def calc_logfile_mu_density_bygantry(index, config, logfile_group,
                                     gantry_angle, grid_resolution=1):
    logfile_mu_density = None

    for filehash in logfile_group:
        filepath = get_filepath(index, config, filehash)
        logfile_delivery_data = DeliveryDatabases.from_logfile(filepath)

        a_logfile_mu_density = [
            get_grid(grid_resolution=grid_resolution),
            logfile_delivery_data.mudensity(
                gantry_angle, grid_resolution=grid_resolution)
        ]

        if logfile_mu_density is None:
            logfile_mu_density = a_logfile_mu_density
        else:
            assert np.all(logfile_mu_density[0] == a_logfile_mu_density[0])
            assert np.all(logfile_mu_density[1] == a_logfile_mu_density[1])
            logfile_mu_density[2] += a_logfile_mu_density[2]

    return logfile_mu_density


def compare_logfile_group_bygantry(index, config, cursor, logfile_group,
                                   gantry_angle, grid_resolution=1):
    field_id = get_field_id_from_logfile_group(index, logfile_group)

    mosaiq_delivery_data = multi_fetch_and_verify_mosaiq(cursor, field_id)

    mosaiq_mu_density = [
        get_grid(grid_resolution=grid_resolution),
        mosaiq_delivery_data.mudensity(
            gantry_angle, grid_resolution=grid_resolution)
    ]

    normalisation = calc_normalisation(mosaiq_delivery_data)

    logfile_mu_density = calc_logfile_mu_density_bygantry(
        index, config, logfile_group, gantry_angle)

    grid_xx = logfile_mu_density[0]
    grid_yy = logfile_mu_density[1]
    assert np.all(grid_xx == mosaiq_mu_density[0])
    assert np.all(grid_yy == mosaiq_mu_density[1])

    comparison = calc_comparison(
        logfile_mu_density[2], mosaiq_mu_density[2], normalisation)

    print(comparison)
    plot_results(
        grid_xx, grid_yy, logfile_mu_density[2], mosaiq_mu_density[2])

    return comparison


# -----------------------------------------------------------------------------
# TODO Refactor the following functions to be in terms of the above.
# -----------------------------------------------------------------------------


def get_logfile_delivery_data_bygantry(index, config, logfile_groups,
                                       mosaiq_gantry_angles):
    logfile_delivery_data_bygantry = dict()

    for logfile_group in logfile_groups:
        logfile_delivery_data_bygantry[logfile_group] = dict()

        for file_hash in logfile_group:
            filepath = get_filepath(index, config, file_hash)
            logfile_delivery_data = DeliveryDatabases.from_logfile(
                filepath)
            mu = np.array(logfile_delivery_data.monitor_units)

            filtered = logfile_delivery_data.filter_cps()

            mu = filtered.monitor_units
            mlc = filtered.mlc
            jaw = filtered.jaw
            logfile_gantry_angles = filtered.gantry

            gantry_tolerance = get_gantry_tolerance(index, file_hash, config)
            unique_logfile_gantry_angles = np.unique(logfile_gantry_angles)

            assert_array_agreement(
                unique_logfile_gantry_angles, mosaiq_gantry_angles,
                gantry_tolerance)

            logfile_delivery_data_bygantry[logfile_group][file_hash] = dict()

            for mosaiq_gantry_angle in mosaiq_gantry_angles:
                logfile_delivery_data_bygantry[logfile_group][file_hash][mosaiq_gantry_angle] = dict(
                )
                agrees_within_tolerance = (
                    np.abs(logfile_gantry_angles - mosaiq_gantry_angle) <= gantry_tolerance)

                logfile_delivery_data_bygantry[logfile_group][file_hash][
                    mosaiq_gantry_angle]['mu'] = mu[agrees_within_tolerance]
                logfile_delivery_data_bygantry[logfile_group][file_hash][
                    mosaiq_gantry_angle]['mlc'] = mlc[agrees_within_tolerance]
                logfile_delivery_data_bygantry[logfile_group][file_hash][
                    mosaiq_gantry_angle]['jaw'] = jaw[agrees_within_tolerance]

    return logfile_delivery_data_bygantry


def get_logfile_mu_density_bygantry(logfile_groups, mosaiq_gantry_angles,
                                    logfile_delivery_data_bygantry):
    logfile_mu_density_bygantry = dict()

    for logfile_group in logfile_groups:
        delivery_data = logfile_delivery_data_bygantry[logfile_group]
        logfile_mu_density_bygantry[logfile_group] = dict()

        for file_hash in logfile_group:
            for mosaiq_gantry_angle in mosaiq_gantry_angles:
                num_control_points = len(
                    delivery_data[file_hash][mosaiq_gantry_angle]['mu'])
                if num_control_points > 0:
                    mu_density = [
                        get_grid(),
                        delivery_data[file_hash][mosaiq_gantry_angle].mudensity()
                    ]

                    if mosaiq_gantry_angle not in logfile_mu_density_bygantry[logfile_group]:
                        logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle] = list(
                            mu_density)
                    else:
                        assert np.all(
                            logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][0] == mu_density[0])
                        assert np.all(
                            logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][1] == mu_density[1])
                        logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][2] += mu_density[2]

    return logfile_mu_density_bygantry


def get_mosaiq_delivery_data_bygantry(mosaiq_delivery_data):
    mu = np.array(mosaiq_delivery_data.monitor_units)
    mlc = np.array(mosaiq_delivery_data.mlc)
    jaw = np.array(mosaiq_delivery_data.jaw)
    gantry_angles = np.array(mosaiq_delivery_data.gantry)
    unique_mosaiq_gantry_angles = np.unique(gantry_angles)

    mosaiq_delivery_data_bygantry = dict()

    for mosaiq_gantry_angle in unique_mosaiq_gantry_angles:
        gantry_angle_matches = gantry_angles == mosaiq_gantry_angle

        diff_mu = np.concatenate([[0], np.diff(mu)])[gantry_angle_matches]
        gantry_angle_specific_mu = np.cumsum(diff_mu)

        mosaiq_delivery_data_bygantry[mosaiq_gantry_angle] = dict()
        mosaiq_delivery_data_bygantry[mosaiq_gantry_angle]['mu'] = gantry_angle_specific_mu
        mosaiq_delivery_data_bygantry[mosaiq_gantry_angle]['mlc'] = mlc[gantry_angle_matches]
        mosaiq_delivery_data_bygantry[mosaiq_gantry_angle]['jaw'] = jaw[gantry_angle_matches]

    return mosaiq_delivery_data_bygantry


def get_mosaiq_mu_density_bygantry(mosaiq_delivery_data_bygantry):
    mosaiq_mu_density_bygantry = dict()
    mosaiq_gantry_angles = mosaiq_delivery_data_bygantry.keys()

    for mosaiq_gantry_angle in mosaiq_gantry_angles:
        mu_density = [
            get_grid(),
            mosaiq_delivery_data_bygantry[mosaiq_gantry_angle].mudensity()
        ]
        mosaiq_mu_density_bygantry[mosaiq_gantry_angle] = mu_density

    return mosaiq_mu_density_bygantry


def get_comparison_results(mosaiq_mu_density_bygantry,
                           logfile_mu_density_bygantry, normalisation):
    comparison_results = dict()
    mosaiq_gantry_angles = mosaiq_mu_density_bygantry.keys()
    logfile_groups = list(logfile_mu_density_bygantry.keys())

    for mosaiq_gantry_angle in mosaiq_gantry_angles:
        comparison_results[mosaiq_gantry_angle] = dict()
        comparison_results[mosaiq_gantry_angle]['comparisons'] = {}

        grid_xx = mosaiq_mu_density_bygantry[mosaiq_gantry_angle][0]
        grid_yy = mosaiq_mu_density_bygantry[mosaiq_gantry_angle][1]
        mosaiq_mu_density = mosaiq_mu_density_bygantry[mosaiq_gantry_angle][2]

        for logfile_group in logfile_groups:
            assert np.all(
                grid_xx == logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][0])
            assert np.all(
                grid_yy == logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][1])

            logfile_mu_density = logfile_mu_density_bygantry[logfile_group][mosaiq_gantry_angle][2]

            comparison = calc_comparison(
                logfile_mu_density, mosaiq_mu_density, normalisation)
            comparison_results[mosaiq_gantry_angle]['comparisons'][logfile_group] = comparison

        comparisons = np.array([
            comparison_results[mosaiq_gantry_angle]['comparisons'][logfile_group]
            for logfile_group in logfile_groups
        ])

        comparison_results[mosaiq_gantry_angle]['median'] = np.median(
            comparisons)
        ref = np.argmin(np.abs(
            comparisons -
            comparison_results[mosaiq_gantry_angle]['median']
        ))
        comparison_results[mosaiq_gantry_angle]['median_filehash_group'] = logfile_groups[ref]
        comparison_results[mosaiq_gantry_angle]['filehash_groups'] = logfile_groups

    return comparison_results


def get_mu_densities_for_file_hashes(index, config, cursor, file_hashes):
    field_ids = {
        index[file_hash]['delivery_details']['field_id']
        for file_hash in file_hashes
    }

    assert len(field_ids) == 1
    field_id = field_ids.pop()

    logfile_groups = group_consecutive_logfiles(file_hashes, index)
    logfile_groups = [
        tuple(group)
        for group in logfile_groups
    ]

    mosaiq_delivery_data = multi_fetch_and_verify_mosaiq(cursor, field_id)
    mosaiq_gantry_angles = np.unique(mosaiq_delivery_data.gantry)

    logfile_delivery_data_bygantry = get_logfile_delivery_data_bygantry(
        index, config, logfile_groups, mosaiq_gantry_angles)
    logfile_mu_density_bygantry = get_logfile_mu_density_bygantry(
        logfile_groups, mosaiq_gantry_angles, logfile_delivery_data_bygantry)
    mosaiq_delivery_data_bygantry = get_mosaiq_delivery_data_bygantry(
        mosaiq_delivery_data)
    mosaiq_mu_density_bygantry = get_mosaiq_mu_density_bygantry(
        mosaiq_delivery_data_bygantry)

    normalisation = calc_normalisation(mosaiq_delivery_data)

    return (
        mosaiq_mu_density_bygantry, logfile_mu_density_bygantry,
        normalisation)


def get_mu_densities_for_field_id(index, config, cursor, field_id, field_id_grouped_hashes):
    file_hashes = np.array(field_id_grouped_hashes[field_id])
    mu_densities = get_mu_densities_for_file_hashes(
        index, config, cursor, file_hashes)

    return mu_densities


def get_comparisons_byfield(index, config, cursor, field_ids,
                            field_id_grouped_hashes):
    comparisons_byfield = dict()

    for field_id in field_ids:
        mu_densities = get_mu_densities_for_field_id(
            index, config, cursor, field_id, field_id_grouped_hashes)
        comparison_results = get_comparison_results(*mu_densities)

        comparisons_byfield[field_id] = comparison_results

    return comparisons_byfield
