# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Codec for processing vulnerability curves
from XML files.

A DOM version of the vulnerability model parser,
that takes into account the really small size of this input file.
"""

from lxml import etree

from openquake import kvs
from openquake import producer
from openquake import shapes
from openquake import xml
from openquake.nrml.utils import nrml_schema_file
from openquake.xml import NRML


def _parse_set_attributes(vulnerability_set):
    """Parse and return the attributes for all the
    vulnerability functions defined in this set of the NRML file."""

    imls = vulnerability_set.find(".//%sIML" % NRML)

    vuln_function = {"IMT": imls.attrib["IMT"]}

    vuln_function["IML"] = \
            [float(x) for x in imls.text.strip().split()]

    vuln_function["vulnerabilitySetID"] = \
            vulnerability_set.attrib["vulnerabilitySetID"]

    vuln_function["assetCategory"] = \
            vulnerability_set.attrib["assetCategory"]

    vuln_function["lossCategory"] = \
            vulnerability_set.attrib["lossCategory"]

    return vuln_function


class VulnerabilityModelFile(producer.FileProducer):
    """This class parsers a vulnerability model NRML file.

    The class is implemented as a generator. For each vulnerability
    function in the parsed instance document it yields a dictionary
    with all the data defined for that function.
    """

    def __init__(self, path):
        producer.FileProducer.__init__(self, path)
        nrml_schema = etree.XMLSchema(etree.parse(nrml_schema_file()))
        self.vuln_model = etree.parse(self.path).getroot()
        if not nrml_schema.validate(self.vuln_model):
            raise xml.XMLValidationError(
                nrml_schema.error_log.last_error, path)
        model_el = self.vuln_model.getchildren()[0]
        if model_el.tag != "%svulnerabilityModel" % NRML:
            raise xml.XMLMismatchError(
                path, 'vulnerabilityModel', str(model_el.tag)[len(NRML):])

    def filter(self, region_constraint=None, attribute_constraint=None):
        """Filtering is not needed/supported for the vulnerability model."""

    def _parse(self):
        """Parse the vulnerability model."""

        for vuln_set in self.vuln_model.findall(
                ".//%sdiscreteVulnerabilitySet" % NRML):

            vuln_function = _parse_set_attributes(vuln_set)

            for raw_vuln_function in vuln_set.findall(
                    ".//%sdiscreteVulnerability" % NRML):

                loss_ratios = [float(x) for x in
                        raw_vuln_function.find(
                        "%slossRatio" % NRML).text.strip().split()]

                coefficients_variation = [float(x) for x in
                        raw_vuln_function.find(
                        "%scoefficientsVariation" % NRML)
                        .text.strip().split()]

                vuln_function["ID"] = \
                        raw_vuln_function.attrib["vulnerabilityFunctionID"]

                vuln_function["probabilisticDistribution"] = \
                        raw_vuln_function.attrib["probabilisticDistribution"]

                vuln_function["lossRatio"] = loss_ratios
                vuln_function["coefficientsVariation"] = coefficients_variation

                yield dict(vuln_function)


def load_vulnerability_model(job_id, path, retrofitted=False):
    """Load and store the vulnerability model defined in the
    given NRML file in the underlying kvs system."""

    vulnerability_model = {}
    parser = VulnerabilityModelFile(path)

    for vuln_curve in parser:

        vuln_func = shapes.VulnerabilityFunction(
            vuln_curve['IML'],
            vuln_curve['lossRatio'],
            vuln_curve['coefficientsVariation'],
            vuln_curve['probabilisticDistribution'])

        vulnerability_model[vuln_curve["ID"]] = vuln_func.to_json()

    kvs.set_value_json_encoded(kvs.tokens.vuln_key(job_id, retrofitted),
            vulnerability_model)


def load_vuln_model_from_kvs(job_id, retrofitted=False):
    """Load the vulnerability model from kvs for the given job."""

    vulnerability_model = kvs.get_value_json_decoded(
            kvs.tokens.vuln_key(job_id, retrofitted))

    vulnerability_curves = {}

    if vulnerability_model is not None:
        for k, v in vulnerability_model.items():
            vulnerability_curves[k] = shapes.VulnerabilityFunction.from_json(v)

    return vulnerability_curves
