# Copyright (c) 2019 University of Illinois and others. All rights reserved.
#
# This program and the accompanying materials are made available under the
# terms of the Mozilla Public License v2.0 which accompanies this distribution,
# and is available at https://www.mozilla.org/en-US/MPL/2.0/


"""Buried Pipeline Damage Analysis with Repair Rate Calculation

"""

import concurrent.futures
import math
from itertools import repeat

from pyincore import BaseAnalysis, HazardService, FragilityService, \
    AnalysisUtil, GeoUtil
from pyincore.analyses.pipelinedamagerepairrate.pipelineutil import \
    PipelineUtil


class PipelineDamageRepairRate(BaseAnalysis):
    """Computes pipeline damage for a hazard.

    Args:
        incore_client: Service client with authentication info

    """

    def __init__(self, incore_client):
        self.hazardsvc = HazardService(incore_client)
        self.fragilitysvc = FragilityService(incore_client)

        super(PipelineDamageRepairRate, self).__init__(incore_client)

    def run(self):
        """Execute pipeline damage analysis """
        # Pipeline dataset
        pipeline_dataset = self.get_input_dataset(
            "pipeline").get_inventory_reader()

        # get input hazard
        hazard, hazard_type, hazard_dataset_id = self.create_hazard_object_from_input_params()

        user_defined_cpu = 1

        if not self.get_parameter("num_cpu") is None and self.get_parameter(
                "num_cpu") > 0:
            user_defined_cpu = self.get_parameter("num_cpu")

        dataset_size = len(pipeline_dataset)
        num_workers = AnalysisUtil.determine_parallelism_locally(self,
                                                                 dataset_size,
                                                                 user_defined_cpu)

        avg_bulk_input_size = int(dataset_size / num_workers)
        inventory_args = []
        count = 0
        inventory_list = list(pipeline_dataset)
        while count < len(inventory_list):
            inventory_args.append(
                inventory_list[count:count + avg_bulk_input_size])
            count += avg_bulk_input_size

        (ds_results, damage_results) = self.pipeline_damage_concurrent_future(
            self.pipeline_damage_analysis_bulk_input, num_workers,
            inventory_args, repeat(hazard), repeat(hazard_type), repeat(hazard_dataset_id))

        self.set_result_csv_data("result", ds_results, name=self.get_parameter("result_name"))
        self.set_result_json_data("metadata",
                                  damage_results,
                                  name=self.get_parameter("result_name") + "_additional_info")

        return True

    def pipeline_damage_concurrent_future(self, function_name, num_workers,
                                          *args):
        """Utilizes concurrent.future module.

        Args:
            function_name (function): The function to be parallelized.
            num_workers (int): Maximum number workers in parallelization.
            *args: All the arguments in order to pass into parameter function_name.

        Returns:
            list: A list of ordered dictionaries with building damage values and other data/metadata.

        """
        output_ds = []
        output_dmg = []
        with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
            for ret1, ret2 in executor.map(function_name, *args):
                output_ds.extend(ret1)
                output_dmg.extend(ret2)

        return output_ds, output_dmg

    def pipeline_damage_analysis_bulk_input(self, pipelines, hazard, hazard_type, hazard_dataset_id):
        """Run pipeline damage analysis for multiple pipelines.

        Args:
            pipelines (list): multiple pipelines from pieline dataset.
            hazard (obj): Hazard object
            hazard_type (str): Hazard type
            hazard_dataset_id (str): An id of the hazard exposure.

        Returns:
            ds_results (list): A list of ordered dictionaries with pipeline damage values and other data/metadata.
            damage_results (list): A list of ordered dictionaries with pipeline damage metadata.
        """
        # Get Fragility key
        fragility_key = self.get_parameter("fragility_key")
        if fragility_key is None:
            fragility_key = PipelineUtil.DEFAULT_TSU_FRAGILITY_KEY if hazard_type == 'tsunami' else \
                PipelineUtil.DEFAULT_EQ_FRAGILITY_KEY
            self.set_parameter("fragility_key", fragility_key)

        # get fragility set
        fragility_sets = self.fragilitysvc.match_inventory(
            self.get_input_dataset("dfr3_mapping_set"), pipelines, fragility_key)

        # Get Liquefaction Fragility Key
        liquefaction_fragility_key = self.get_parameter(
            "liquefaction_fragility_key")
        if hazard_type == "earthquake" and liquefaction_fragility_key is None:
            liquefaction_fragility_key = PipelineUtil.LIQ_FRAGILITY_KEY

        # Liquefaction
        use_liquefaction = False
        if hazard_type == "earthquake" and self.get_parameter(
                "use_liquefaction") is not None:
            use_liquefaction = self.get_parameter("use_liquefaction")

        # Get geology dataset id
        geology_dataset_id = self.get_parameter(
            "liquefaction_geology_dataset_id")
        fragility_sets_liq = None
        if geology_dataset_id is not None:
            fragility_sets_liq = self.fragilitysvc.match_inventory(
                self.get_input_dataset("dfr3_mapping_set"), pipelines,
                liquefaction_fragility_key)

        values_payload = []
        values_payload_liq = []  # for liquefaction if used
        unmapped_pipelines = []
        mapped_pipelines = []
        for pipeline in pipelines:
            # if find a match fragility for that pipeline
            if pipeline["id"] in fragility_sets.keys():
                fragility_set = fragility_sets[pipeline["id"]]
                location = GeoUtil.get_location(pipeline)
                loc = str(location.y) + "," + str(location.x)
                demands = fragility_set.demand_types
                units = fragility_set.demand_units
                value = {
                    "demands": demands,
                    "units": units,
                    "loc": loc
                }
                values_payload.append(value)
                mapped_pipelines.append(pipeline)

                # Check if liquefaction is applicable
                if use_liquefaction and \
                        geology_dataset_id is not None and \
                        fragility_sets_liq is not None and \
                        pipeline["id"] in fragility_sets_liq:
                    fragility_set_liq = fragility_sets_liq[pipeline["id"]]
                    demands_liq = fragility_set_liq.demand_types
                    units_liq = fragility_set_liq.demand_units
                    value_liq = {
                        "demands": demands_liq,
                        "units": units_liq,
                        "loc": loc
                    }
                    values_payload_liq.append(value_liq)
            else:
                unmapped_pipelines.append(pipeline)
        del pipelines

        if hazard_type == 'earthquake' or hazard_type == 'tsunami':
            hazard_resp = hazard.read_hazard_values(values_payload, self.hazardsvc)
        else:
            raise ValueError("The provided hazard type is not supported yet by this analysis")

        liquefaction_resp = None
        # Check if liquefaction is applicable
        if use_liquefaction is True and \
                fragility_sets_liq is not None and \
                geology_dataset_id is not None:
            liquefaction_resp = self.hazardsvc.post_liquefaction_values(hazard_dataset_id, geology_dataset_id,
                                                                        values_payload_liq)

        # calculate LS and DS
        ds_results = []
        damage_results = []
        for i, pipeline in enumerate(mapped_pipelines):
            # default
            pgv_repairs = None
            pgd_repairs = 0.0
            total_repair_rate = None
            break_rate = None
            leak_rate = None
            failure_probability = None
            num_pgv_repairs = None
            num_pgd_repairs = 0.0
            num_repairs = None

            liq_hazard_vals = None
            liq_demand_types = None
            liq_demand_units = None
            liquefaction_prob = None

            ds_result = dict()
            damage_result = dict()
            ds_result['guid'] = pipeline['properties']['guid']
            damage_result['guid'] = pipeline['properties']['guid']

            fragility_set = fragility_sets[pipeline["id"]]
            # TODO assume there is only one curve
            fragility_curve = fragility_set.fragility_curves[0]

            hazard_vals = AnalysisUtil.update_precision_of_lists(hazard_resp[i]["hazardValues"])
            demand_types = hazard_resp[i]["demands"]
            demand_units = hazard_resp[i]["units"]

            hval_dict = dict()
            for j, d in enumerate(fragility_set.demand_types):
                hval_dict[d] = hazard_vals[j]

            if not AnalysisUtil.do_hazard_values_have_errors(hazard_resp[i]["hazardValues"]):
                pipeline_args = fragility_set.construct_expression_args_from_inventory(pipeline)
                pgv_repairs = \
                    fragility_curve.solve_curve_expression(
                        hval_dict, fragility_set.curve_parameters, **pipeline_args)
                # Convert PGV repairs to SI units
                pgv_repairs = PipelineUtil.convert_result_unit(fragility_curve.return_type["unit"], pgv_repairs)

                length = PipelineUtil.get_pipe_length(pipeline)

                # Number of PGV repairs
                num_pgv_repairs = pgv_repairs * length

                # Check if liquefaction is applicable
                if use_liquefaction is True \
                        and fragility_sets_liq is not None \
                        and geology_dataset_id is not None \
                        and liquefaction_resp is not None:
                    fragility_set_liq = fragility_sets_liq[pipeline["id"]]

                    # TODO assume there is only one curve
                    liq_fragility_curve = fragility_set_liq.fragility_curves[0]

                    liq_hazard_vals = AnalysisUtil.update_precision_of_lists(liquefaction_resp[i]["pgdValues"])
                    liq_demand_types = liquefaction_resp[i]["demands"]
                    liq_demand_units = liquefaction_resp[i]["units"]
                    liquefaction_prob = liquefaction_resp[i]['liqProbability']
                    liq_hval_dict = dict()
                    for j, d in enumerate(liquefaction_resp[i]["demands"]):
                        liq_hval_dict[d] = liq_hazard_vals[j]

                    # !important! removing the liqProbability and passing in the "diameter"
                    # no fragility is actually using liqProbability
                    pipeline_args = fragility_set_liq.construct_expression_args_from_inventory(pipeline)
                    pgd_repairs = \
                        liq_fragility_curve.solve_curve_expression(
                            liq_hval_dict, fragility_set_liq.curve_parameters, **pipeline_args)
                    # Convert PGD repairs to SI units
                    pgd_repairs = PipelineUtil.convert_result_unit(liq_fragility_curve.return_type["unit"], pgd_repairs)
                    num_pgd_repairs = pgd_repairs * length

                    # record results
                    if 'pipetype' in pipeline['properties']:
                        damage_result['pipeclass'] = pipeline['properties']['pipetype']
                    elif 'pipelinesc' in pipeline['properties']:
                        damage_result['pipeclass'] = pipeline['properties']['pipelinesc']
                    else:
                        damage_result['pipeclass'] = ""

                break_rate = 0.2 * pgv_repairs + 0.8 * pgd_repairs
                leak_rate = 0.8 * pgv_repairs + 0.2 * pgd_repairs
                total_repair_rate = pgd_repairs + pgv_repairs
                failure_probability = 1 - math.exp(-1.0 * break_rate * length)
                num_repairs = num_pgd_repairs + num_pgv_repairs

            ds_result['pgvrepairs'] = pgv_repairs
            ds_result['pgdrepairs'] = pgd_repairs
            ds_result['repairspkm'] = total_repair_rate
            ds_result['breakrate'] = break_rate
            ds_result['leakrate'] = leak_rate
            ds_result['failprob'] = failure_probability
            ds_result['numpgvrpr'] = num_pgv_repairs
            ds_result['numpgdrpr'] = num_pgd_repairs
            ds_result['numrepairs'] = num_repairs
            ds_result['haz_expose'] = AnalysisUtil.get_exposure_from_hazard_values(hazard_vals, hazard_type)

            damage_result['fragility_id'] = fragility_set.id
            damage_result['demandtypes'] = demand_types
            damage_result['demandunits'] = demand_units
            damage_result['hazardtype'] = hazard_type
            damage_result['hazardval'] = hazard_vals

            # Check if liquefaction is applicable
            if use_liquefaction is True \
                    and fragility_sets_liq is not None \
                    and geology_dataset_id is not None:
                damage_result['liq_fragility_id'] = fragility_sets_liq[pipeline["id"]].id
                damage_result['liqdemandtypes'] = liq_demand_types
                damage_result['liqdemandunits'] = liq_demand_units
                damage_result['liqhazval'] = liq_hazard_vals
                damage_result['liqprobability'] = liquefaction_prob
            else:
                damage_result['liq_fragility_id'] = None
                damage_result['liqdemandtypes'] = None
                damage_result['liqdemandunits'] = None
                damage_result['liqhazval'] = None
                damage_result['liqprobability'] = None

            ds_results.append(ds_result)
            damage_results.append(damage_result)

        # pipelines do not have matched mappings
        for pipeline in unmapped_pipelines:
            ds_result = dict()
            ds_result['guid'] = pipeline['properties']['guid']

            damage_result = dict()
            damage_result['guid'] = pipeline['properties']['guid']
            if 'pipetype' in pipeline['properties']:
                damage_result['pipeclass'] = pipeline['properties']['pipetype']
            elif 'pipelinesc' in pipeline['properties']:
                damage_result['pipeclass'] = pipeline['properties']['pipelinesc']
            else:
                damage_result['pipeclass'] = ""

            damage_result['fragility_id'] = None
            damage_result['demandtypes'] = None
            damage_result['demandunits'] = None
            damage_result['hazardtype'] = None
            damage_result['hazardval'] = None
            damage_result['liq_fragility_id'] = None
            damage_result['liqdemandtypes'] = None
            damage_result['liqdemandunits'] = None
            damage_result['liqhazval'] = None
            damage_result['liqhazval'] = None

            ds_results.append(ds_result)
            damage_results.append(damage_result)

        return ds_results, damage_results

    def get_spec(self):
        """Get specifications of the pipeline damage analysis.

        Returns:
            obj: A JSON object of specifications of the pipeline damage analysis.

        """
        return {
            'name': 'pipeline-damage',
            'description': 'buried pipeline damage analysis',
            'input_parameters': [
                {
                    'id': 'result_name',
                    'required': True,
                    'description': 'result dataset name',
                    'type': str
                },
                {
                    'id': 'hazard_type',
                    'required': False,
                    'description': 'Hazard Type (e.g. earthquake)',
                    'type': str
                },
                {
                    'id': 'hazard_id',
                    'required': False,
                    'description': 'Hazard ID',
                    'type': str
                },
                {
                    'id': 'fragility_key',
                    'required': False,
                    'description': 'Fragility key to use in mapping dataset',
                    'type': str
                },
                {
                    'id': 'use_liquefaction',
                    'required': False,
                    'description': 'Use liquefaction',
                    'type': bool
                },
                {
                    'id': 'liquefaction_fragility_key',
                    'required': False,
                    'description': 'Fragility key to use in liquefaction mapping dataset',
                    'type': str
                },
                {
                    'id': 'num_cpu',
                    'required': False,
                    'description': 'If using parallel execution, the number of cpus to request',
                    'type': int
                },
                {
                    'id': 'liquefaction_geology_dataset_id',
                    'required': False,
                    'description': 'Geology dataset id',
                    'type': str,
                }
            ],
            'input_hazards': [
                {
                    'id': 'hazard',
                    'required': False,
                    'description': 'Hazard object',
                    'type': ["earthquake", "tsunami"]
                },
            ],
            'input_datasets': [
                {
                    'id': 'pipeline',
                    'required': True,
                    'description': 'Pipeline Inventory',
                    'type': ['ergo:buriedPipelineTopology', 'ergo:pipeline'],
                },
                {
                    'id': 'dfr3_mapping_set',
                    'required': True,
                    'description': 'DFR3 Mapping Set Object',
                    'type': ['incore:dfr3MappingSet'],
                }
            ],
            'output_datasets': [
                {
                    'id': 'result',
                    'parent_type': 'pipeline',
                    'type': 'ergo:pipelineDamageVer3'
                },
                {
                    'id': 'metadata',
                    'parent_type': 'pipeline',
                    'description': 'additional metadata in json file about applied hazard value and '
                                   'fragility',
                    'type': 'incore:pipelineDamageSupplement'
                }
            ]
        }
