#!/usr/bin/env python

import os
import argparse
import warnings
import pkg_resources

import numpy as np
from collections import OrderedDict
from configobj import ConfigObj

from shakemap.utils.config import get_custom_validator
from shakemap.utils.config import config_error

#-------------------------------------------------------------------------------
# Openquake utilities
from openquake.hazardlib import imt, const

from mapio.geodict import GeoDict
from mapio.gmt import GMTGrid
from mapio.shake import ShakeGrid

# Shakemap imports
from shakelib.rupture.point_rupture import PointRupture
from shakelib.rupture.factory import get_rupture
from shakelib.rupture.origin import Origin
from shakelib.distance import Distance
from shakelib.sites import Sites
from shakelib.multigmpe import MultiGMPE
from shakelib.directivity.rowshandel2013 import Rowshandel2013
from datetime import datetime
from shakelib.gmice.wgrw12 import WGRW12
from shakelib.virtualipe import VirtualIPE

from impactutils.io.cmd import get_command_output

from scenarios.utils import get_extent


def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")


def main(args):
    id_str = args.event

    spec_file = pkg_resources.resource_filename(
        'scenarios', os.path.join('data', 'configspec.conf'))
    validator = get_custom_validator()
    config = ConfigObj(os.path.join(os.path.expanduser('~'), 'scenarios.conf'),
                       configspec=spec_file)

    shakehome = config['system']['shakehome']

    datdir = os.path.join(shakehome, 'data')
    evt_dir = os.path.join(datdir, id_str)
    input_dir = os.path.join(evt_dir, 'input')
    xml_file = os.path.join(input_dir, 'event.xml')

    #---------------------------------------------------------------------------
    # Read in event.xml and create Origin object
    #---------------------------------------------------------------------------
    origin = Origin.fromFile(xml_file)

    #---------------------------------------------------------------------------
    # Read in rupture
    #---------------------------------------------------------------------------
    cmd = 'ls %s/*rupture.json' % (input_dir)
    rc, so, se = get_command_output(cmd)
    if rc is False:
        cmd = 'ls %s/*_fault.txt' % (input_dir)
        rc, so, se = get_command_output(cmd)

    if args.verbose is True:
        print('Rupture file: %s\n' % so.decode('utf-8').strip())

    if rc is True:
        # There is a rupture
        ruptfile = so.decode('utf-8').strip()
        rupt = get_rupture(origin, ruptfile)
    else:
        rupt = PointRupture(origin)

    # Set the dx for the rupture meshing
    rupt._mesh_dx = args.mesh_dx

    #---------------------------------------------------------------------------
    # Construct the MultiGMPE, not specific/filtered to an IMT
    #---------------------------------------------------------------------------

    tmp = pkg_resources.resource_filename(
        'scenarios', os.path.join('..', 'data', 'gmpe_sets.conf'))
    config.merge(ConfigObj(tmp, configspec=spec_file))
    tmp = pkg_resources.resource_filename(
        'scenarios', os.path.join('..', 'data', 'modules.conf'))
    config.merge(ConfigObj(tmp, configspec=spec_file))
    results = config.validate(validator)
    if results != True:
        config_error(config, results)

    config = config.dict()

    gmpe = MultiGMPE.from_config(config, verbose=args.verbose)

    #---------------------------------------------------------------------------
    # Compute extent, or get from args
    #---------------------------------------------------------------------------
    if args.extent is None:
        lonmin, lonmax, latmin, latmax = get_extent(origin, rupt)
    else:
        lonmin, latmin, lonmax, latmax = args.extent

    # Adjust extent to be divisible by resolution
    res = args.res
    lonmin = res * np.round(lonmin / res)
    lonmax = res * np.round(lonmax / res)
    latmin = res * np.round(latmin / res)
    latmax = res * np.round(latmax / res)

    lonspan = float(lonmax - lonmin)
    latspan = float(latmax - latmin)

    # Adjust number of cells if necessary
    nx = np.floor(lonspan / res) + 1
    ny = np.floor(latspan / res) + 1
    ncell = nx * ny
    nmax = args.max
    if ncell > nmax:
        res = (-(latspan + lonspan) -
               np.sqrt(latspan**2 + lonspan**2 + 2 * latspan * lonspan *
                       (2 * nmax - 1))) / (2 * (1 - nmax))
        warnings.warn(
            'resolution adjusted due to max number of cells allowed.')
        nx = np.floor(lonspan / res) + 1
        ny = np.floor(latspan / res) + 1
        ncell = nx * ny

    # Geodictionary for this ShakeMap
    tmpdict = {'xmin': lonmin, 'xmax': lonmax,
               'ymin': latmin, 'ymax': latmax,
               'dx': res, 'dy': res,
               'nx': nx, 'ny': ny}
    smdict = GeoDict(tmpdict, adjust='bounds')

    if args.verbose is True:
        print('Geodictionary:')
        print(smdict)
        print('')

    #---------------------------------------------------------------------------
    # Make rupture context
    #---------------------------------------------------------------------------
    rx = rupt.getRuptureContext(gmpe)

    if args.verbose is True:
        print('Rupture context:')
        print('Mag: %s' % rx.mag)
        print('Hyp lat: %s' % rx.hypo_lat)
        print('Hyp lon: %s' % rx.hypo_lon)
        print('Hyp dep: %s\n' % rx.hypo_depth)

    #---------------------------------------------------------------------------
    # Vs30 stuff
    #---------------------------------------------------------------------------
    vs30filename = config['data']['vs30file']
    vs30grid = GMTGrid.load(vs30filename, smdict, resample=True)

    # Sites object
    sites = Sites(vs30grid)
    sx = sites.getSitesContext()
    sx_rock = sites.getSitesContext(rock_vs30=760)

    # Clip Vs30 do avoid interpolation error
    sx.vs30 = np.clip(sx.vs30, 0, 2000)

    if args.verbose is True:
        print('Sites context:')
        print('Lons: %s to %s' % (np.min(sx.lons), np.max(sx.lons)))
        print('Lats: %s to %s' % (np.min(sx.lats), np.max(sx.lats)))
        print('Min Vs30: %s' % np.min(sx.vs30))
        print('Max Vs30: %s\n' % np.max(sx.vs30))

    #---------------------------------------------------------------------------
    # Standard deviation stuff
    #---------------------------------------------------------------------------

    # Only use total standard deviation for scenarios since
    # we never have data to get bias.
    stddev_types = [const.StdDev.TOTAL]

    #---------------------------------------------------------------------------
    # Intensity measures (excluding MI)
    #---------------------------------------------------------------------------

    # Mapping between the IM notation in ShakeMap and the
    # OpenQuake notation for the IMs taht we want
    imt_dict = {'pga': 'PGA', 'pgv': 'PGV', 'psa03': 'SA(0.3)',
                'psa10': 'SA(1.0)', 'psa30': 'SA(3.0)'}

    #---------------------------------------------------------------------------
    # Mesh calculations
    #---------------------------------------------------------------------------
    lats = np.linspace(smdict.ymax, smdict.ymin, smdict.ny)
    lons = np.linspace(smdict.xmin, smdict.xmax, smdict.nx)
    lon, lat = np.meshgrid(lons, lats)
    dep = np.zeros_like(lon)

    if args.verbose is True:
        print('Mesh:')
        print('Lons: %s to %s' % (np.min(lons), np.max(lons)))
        print('Lats: %s to %s' % (np.min(lats), np.max(lats)))
        print('mesh_dx: %f\n' % rupt._mesh_dx)

    # Compute distances and site parameters on mesh.
    dist = Distance(gmpe, lon, lat, dep, rupt)
    dx = dist.getDistanceContext()

    if args.verbose is True:
        print('Distance context:')
        print('Min Rrup: %s' % np.min(dx.rrup))
        print('Max Rrup: %s\n' % np.max(dx.rrup))

    #---------------------------------------------------------------------------
    # Intensity measure calculation
    #---------------------------------------------------------------------------

    # Make a dictionary to store intensity measure(s) and
    # their sigmas.

    orig_shape = sx.vs30.shape
    imdict = {'pga': {'mean': np.zeros(orig_shape),
                      'sigma': np.zeros(orig_shape)},
              'pgv': {'mean': np.zeros(orig_shape),
                      'sigma': np.zeros(orig_shape)},
              'psa03': {'mean': np.zeros(orig_shape),
                        'sigma': np.zeros(orig_shape)},
              'psa10': {'mean': np.zeros(orig_shape),
                        'sigma': np.zeros(orig_shape)},
              'psa30': {'mean': np.zeros(orig_shape),
                        'sigma': np.zeros(orig_shape)}}

    # For rock
    imdict_rock = {'pga': {'mean': np.zeros(orig_shape),
                           'sigma': np.zeros(orig_shape)},
                   'pgv': {'mean': np.zeros(orig_shape),
                           'sigma': np.zeros(orig_shape)},
                   'psa03': {'mean': np.zeros(orig_shape),
                             'sigma': np.zeros(orig_shape)},
                   'psa10': {'mean': np.zeros(orig_shape),
                             'sigma': np.zeros(orig_shape)},
                   'psa30': {'mean': np.zeros(orig_shape),
                             'sigma': np.zeros(orig_shape)}}

    #---------------------------------------------------------------------------
    # Directivity
    #---------------------------------------------------------------------------
    if isinstance(origin.directivity, str):
        dirbool = str2bool(origin.directivity)
    else:
        dirbool = origin.directivity
    if dirbool is True:
        R13 = Rowshandel2013.fromSites(
            origin, rupt, sites, dx=1.0, T=[1.0, 3.0],
            a_weight=0.5, mtype=1)
        fd1 = R13.getFd()[0]
        fd3 = R13.getFd()[1]

    #---------------------------------------------------------------------------
    # Evaluate GMPEs
    #---------------------------------------------------------------------------
    for key, val in imt_dict.items():
        iimt = imt.from_string(val)
        if args.verbose is True:
            print("IMT: %s" % val)

        #-----------------------------------------------------------------------
        # Re-construct the MultiGMPE, but specifcy the IMT so that it filters
        # out GMPEs that cannot be evaluated at this IMT
        #-----------------------------------------------------------------------
        gmpe = MultiGMPE.from_config(config, filter_imt=iimt,
                                     verbose=args.verbose)
        lnmu, lnsd = gmpe.get_mean_and_stddevs(
            sx, rx, dx, iimt, stddev_types)
        lnmu_rock, lnsd_rock = gmpe.get_mean_and_stddevs(
            sx_rock, rx, dx, iimt, stddev_types)

        #-----------------------------------------------------------------------
        # Handle directivity factors
        # NOTE: currently, the Rowshandel model does not provide
        #       equations for adjusting sigma. Asssuming these are
        #       eventually available, need to move the sigma
        #       adjustment into this if-statement.
        #-----------------------------------------------------------------------
        if dirbool is True:
            if (key == 'pgv') | (key == 'psa10'):
                fd = fd1
            elif (key == 'psa30'):
                fd = fd3
            else:
                # no directivity for pga and psa03
                fd = 0

            lnmu = lnmu + fd
            lnmu_rock = lnmu_rock + fd

        # Put into intensity measure dictionary
        imdict[key]['mean'] = lnmu
        imdict[key]['sigma'] = lnsd[0]

        imdict_rock[key]['mean'] = lnmu_rock
        imdict_rock[key]['sigma'] = lnsd_rock[0]

    #-----------------------------------------------------------------------
    # Write files
    #-----------------------------------------------------------------------

    # Loop over intensity dictionary (PGA PGV, PSA03, PSA10, PSA30)
    for key, val in imdict.items():
        if key != 'pgv':
            # Note that the output is in units of ln(g), whereas
            # ShakeMap wants %g
            mgrid = GMTGrid(100 * np.exp(imdict[key]['mean']), smdict)
            sgrid = GMTGrid(imdict[key]['sigma'], smdict)
        else:
            mgrid = GMTGrid(np.exp(imdict[key]['mean']), smdict)
            sgrid = GMTGrid(imdict[key]['sigma'], smdict)

        if args.verbose is True:
            print('Min %s: %s' % (key, np.min(mgrid.getData())))
            print('Max %s: %s\n' % (key, np.max(mgrid.getData())))

        # Write to file
        mgrid.save(os.path.join(input_dir, key + '_estimates.grd'))
        sgrid.save(os.path.join(input_dir, key + '_sd.grd'))

    # Also write directivity factors to a file
    if dirbool is True:
        fd1grd = GMTGrid(fd1, smdict)
        fd3grd = GMTGrid(fd3, smdict)
        fd1grd.save(os.path.join(input_dir, 'fd1.grd'))
        fd3grd.save(os.path.join(input_dir, 'fd3.grd'))

    #---------------------------------------------------------------------------
    # MMI - Use VirtualIPE
    #---------------------------------------------------------------------------
    gmpe = MultiGMPE.from_config(config, filter_imt=imt.PGV(),
                                 verbose=args.verbose)
    gmice = WGRW12()
    vipe = VirtualIPE.fromFuncs(gmpe, gmice)
    if dirbool is True:
        mmi, mmi_sd = vipe.get_mean_and_stddevs(
            sx, rx, dx, imt.MMI(), stddev_types, fd1)
    else:
        mmi, mmi_sd = vipe.get_mean_and_stddevs(
            sx, rx, dx, imt.MMI(), stddev_types)

    mgrid = GMTGrid(mmi, smdict)
    sgrid = GMTGrid(mmi_sd[0], smdict)

    if args.verbose is True:
        print('Min MI: %s' % np.min(mgrid.getData()))
        print('Max MI: %s\n' % np.max(mgrid.getData()))

    # Write to file
    mgrid.save(os.path.join(input_dir, 'mi_estimates.grd'))
    sgrid.save(os.path.join(input_dir, 'mi_sd.grd'))

    # Write GMPE set name to a file to put into info.json later
    gmpefile = open(os.path.join(input_dir, "gmpe_set_name.txt"), "w")
    gmpefile.write(gmpe.DESCRIPTION)
    gmpefile.close()

    # Need to write rock_grid.xml
    layers = OrderedDict()
    layers['pga'] = 100 * np.exp(imdict_rock['pga']['mean'])
    layers['pgv'] = np.exp(imdict_rock['pgv']['mean'])
    layers['mmi'] = mmi
    layers['psa03'] = 100 * np.exp(imdict_rock['psa03']['mean'])
    layers['psa10'] = 100 * np.exp(imdict_rock['psa10']['mean'])
    layers['psa30'] = 100 * np.exp(imdict_rock['psa30']['mean'])
    shakeDict = {'event_id': id_str,
                 'shakemap_id': id_str,
                 'shakemap_version': 1,
                 'code_version': '4.0',
                 'process_timestamp': datetime.utcnow(),
                 'shakemap_originator': 'us',
                 'map_status': 'RELEASED',
                 'shakemap_event_type': 'SCENARIO'}
    eventDict = {'event_id': id_str,
                 'magnitude': rx.mag,
                 'depth': rx.hypo_depth,
                 'lat': rx.hypo_lat,
                 'lon': rx.hypo_lon,
                 'event_timestamp': datetime.utcnow(),
                 'event_network': 'us',
                 'event_description': ""}
    #---------------------------------------------------------------------------
    # Note: cannot put standard deviations here becaus they go
    # into uncertainty.xml; in the current version of shakemap there
    # is no where to put rock grid standard deviations. We should
    # change this in the next udpate. The uncertainty here is based on
    # instrumental records and so it is not really meaningful for scenarios.
    #---------------------------------------------------------------------------
    uncDict = {'pga': (0, 0),
               'pgv': (0, 0),
               'mmi': (0, 0),
               'psa03': (0, 0),
               'psa10': (0, 0),
               'psa30': (0, 0)}
    shake = ShakeGrid(layers, smdict, eventDict, shakeDict, uncDict)
    shake.save(os.path.join(input_dir, "rock_grid.xml"), version=1)


if __name__ == '__main__':
    desc = '''
    Create the ShakeMap *_estimates.grd and *_sd.grd files. Requires an input
    directory with event.xml and rupture files; the inputs should be created by
    the 'mkinputdir' script or similar.
    '''
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument(
        '-e', '--event', required=True,
        help='Specifies the id of the event to process.')
    parser.add_argument(
        '-r', '--res', default=30 / 60 / 60, type=float,
        help='The resolution in decimal degrees; default is 30/60/60.')
    parser.add_argument(
        '-m', '--max', default=500000, type=int,
        help='Maximum number of cells allowed; '
        'resolution is adjusted to ensure this number is not exceeded; '
        'default is 500,000.')
    parser.add_argument(
        '--mesh_dx', default=0.5, type=float,
        help='The resolution for rupture mesh in km; only used for EdgeRuptures; '
             'default is 0.5.')
    parser.add_argument(
        '--extent', nargs='+', help='Extent: lonmin, latmin, lonmax, latmax.',
        required=False, default=None, type=float)
    parser.add_argument(
        '-v', '--verbose', action="store_true", default=False,
        help='Add verbose output.')
    args = parser.parse_args()
    main(args)
