#!/usr/bin/env python3
"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
The FID metric calculates the distance between two distributions of images.
Typically, we have summary statistics (mean & covariance matrix) of one
of these distributions, while the 2nd distribution is given by a GAN.
When run as a stand-alone program, it compares the distribution of
images that are stored as PNG/JPEG at a specified location with a
distribution given by summary statistics (in pickle format).
The FID is calculated by assuming that X_1 and X_2 are the activations of
the pool_3 layer of the inception net for generated samples and real world
samples respectively.
See --help to see further details.
Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
of Tensorflow
Copyright 2018 Institute of Bioinformatics, JKU Linz
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
   http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import pathlib

import numpy as np
import torch
from scipy import linalg
from torch.nn.functional import adaptive_avg_pool2d
from src.utils import read_config
from src.Models.models import ImitateJoint
import matplotlib.pyplot as plt
from globals import device
from autoencoder import Autoencoder
from PIL import Image

try:
    from tqdm import tqdm
except ImportError:
    # If not tqdm is not available, provide a mock version of it
    def tqdm(x): return x

from src.Models.models import Encoder
from src.utils.generators.wake_sleep_gen import WakeSleepGen
from src.utils.generators.shapenet_generater import Generator
from globals import device

def get_activations(generator, model, batch_size=50, dims=32, verbose=False):
    """Calculates the activations of the pool_3 layer for all images.
    Params:
    -- files       : List of image files paths
    -- model       : Instance of inception model
    -- batch_size  : Batch size of images for the model to process at once.
                     Make sure that the number of samples is a multiple of
                     the batch size, otherwise some samples are ignored. This
                     behavior is retained to match the original FID score
                     implementation.
    -- dims        : Dimensionality of features returned by Inception
    -- cuda        : If set to True, use GPU
    -- verbose     : If set to True and parameter out_step is given, the number
                     of calculated batches is reported.
    Returns:
    -- A numpy array of dimension (num images, dims) that contains the
       activations of the given tensor when feeding inception with the
       query tensor.
    """
    model.eval()

    test_size = 3000

    pred_arr = np.empty((test_size, dims))

    n_batches = test_size // batch_size

    for i in tqdm(range(0, test_size, batch_size)):
        if verbose:
            print('\rPropagating batch %d/%d' % ((i + 1) // batch_size, n_batches),
                  end='', flush=True)
        start = i
        end = i + batch_size

        images = next(generator)
        images = torch.from_numpy(images).to(device)
        if len(images.shape) == 3: # generated samples
            # images = torch.from_numpy(np.random.randint(0, 2, (100, 64, 64))).to(device).float()
            pred = model.encode(images.unsqueeze(1).float())
        else: # cad data
            pred = model.encode(images[-1, :, 0:1, :, :])

        pred = pred.reshape((batch_size,-1, 1, 1))
        # If model output is not scalar, apply global spatial average pooling.
        # This happens if you choose a dimensionality not equal 2048.
        # if pred.size(2) != 1 or pred.size(3) != 1:
        #     pred = adaptive_avg_pool2d(pred, output_size=(1, 1))

        pred_arr[start:end] = pred.cpu().data.numpy().reshape(pred.size(0), -1)

    if verbose:
        print(' done')

    return pred_arr

def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
    """Numpy implementation of the Frechet Distance.
    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
    and X_2 ~ N(mu_2, C_2) is
            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
    Stable version by Dougal J. Sutherland.
    Params:
    -- mu1   : Numpy array containing the activations of a layer of the
               inception net (like returned by the function 'get_predictions')
               for generated samples.
    -- mu2   : The sample mean over activations, precalculated on an
               representative data set.
    -- sigma1: The covariance matrix over activations for generated samples.
    -- sigma2: The covariance matrix over activations, precalculated on an
               representative data set.
    Returns:
    --   : The Frechet Distance.
    """

    mu1 = np.atleast_1d(mu1)
    mu2 = np.atleast_1d(mu2)

    sigma1 = np.atleast_2d(sigma1)
    sigma2 = np.atleast_2d(sigma2)

    assert mu1.shape == mu2.shape, \
        'Training and test mean vectors have different lengths'
    assert sigma1.shape == sigma2.shape, \
        'Training and test covariances have different dimensions'

    diff = mu1 - mu2

    # Product might be almost singular
    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
    if not np.isfinite(covmean).all():
        msg = ('fid calculation produces singular product; '
               'adding %s to diagonal of cov estimates') % eps
        print(msg)
        offset = np.eye(sigma1.shape[0]) * eps
        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

    # Numerical error might give slight imaginary component
    if np.iscomplexobj(covmean):
        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
            m = np.max(np.abs(covmean.imag))
            raise ValueError('Imaginary component {}'.format(m))
        covmean = covmean.real

    tr_covmean = np.trace(covmean)

    return (diff.dot(diff) + np.trace(sigma1) +
            np.trace(sigma2) - 2 * tr_covmean)

def calculate_activation_statistics(files, model, batch_size=50,
                                    dims=32, verbose=False):
    """Calculation of the statistics used by the FID.
    Params:
    -- files       : List of image files paths
    -- model       : Instance of inception model
    -- batch_size  : The images numpy array is split into batches with
                     batch size batch_size. A reasonable batch size
                     depends on the hardware.
    -- dims        : Dimensionality of features returned by Inception
    -- cuda        : If set to True, use GPU
    -- verbose     : If set to True and parameter out_step is given, the
                     number of calculated batches is reported.
    Returns:
    -- mu    : The mean over samples of the activations of the pool_3 layer of
               the inception model.
    -- sigma : The covariance matrix of the activations of the pool_3 layer of
               the inception model.
    """
    act = get_activations(files, model, batch_size, dims, verbose)
    mu = np.mean(act, axis=0)
    sigma = np.cov(act, rowvar=False)
    return mu, sigma

def _compute_statistics_of_path(path, model, batch_size, dims):
    if path.endswith('.npz'):
        f = np.load(path)
        m, s = f['mu'][:], f['sigma'][:]
        f.close()
    else:
        m, s = calculate_activation_statistics(path, model, batch_size,
                                               dims)

    return m, s

def calculate_fid_given_paths(model, images_path, model_path, batch_size, dims=32):
    """Calculates the FID of two paths"""
    if not os.path.exists(images_path):
        raise RuntimeError('Invalid path: %s' % images_path)
    if not os.path.exists(model_path):
        raise RuntimeError('Invalid path: %s' % model_path)

    generator = FidGen(images_path).get_test_data()
    # generator2 = FidGen(model_path).get_test_data()
    cad_generator = Generator().test_gen(batch_size=batch_size,
                                        path="data/cad/cad.h5",
                                        if_augment=False)
    # cad_generator2 = Generator().val_gen(batch_size=batch_size,
    #                                     path="data/cad/cad.h5",
    #                                     if_augment=False)

    m1, s1 = calculate_activation_statistics(cad_generator, model, batch_size,
                                         dims)
    m2, s2 = calculate_activation_statistics(generator, model, batch_size,
                                         dims)

    fid_value = calculate_frechet_distance(m1, s1, m2, s2)

    return fid_value

def get_csgnet():
    config = read_config.Config("config_synthetic.yml")

    # Encoder
    encoder_net = Encoder(config.encoder_drop)
    encoder_net = encoder_net.to(device)

    imitate_net = ImitateJoint(
        hd_sz=config.hidden_size,
        input_size=config.input_size,
        encoder=encoder_net,
        mode=config.mode,
        num_draws=400,
        canvas_shape=config.canvas_shape)
    imitate_net = imitate_net.to(device)

    print("pre loading model")
    pretrained_dict = torch.load(config.pretrain_modelpath, map_location=device)
    imitate_net_dict = imitate_net.state_dict()
    pretrained_dict = {
        k: v
        for k, v in pretrained_dict.items() if k in imitate_net_dict
    }
    imitate_net_dict.update(pretrained_dict)
    imitate_net.load_state_dict(imitate_net_dict)

    return imitate_net.encoder

class FidGen:
    def __init__(self, images_path):
        self.images = torch.load(images_path)
    def get_test_data(self):
        while True:
            for i in range(0, 3000, 100):
                batch_images = self.images[i:i+100]
                yield batch_images

if __name__ == '__main__':
    # model = Autoencoder().to(device)
    # model.load_state_dict(torch.load("trained_models/fid-model2.pth"))
    model = get_csgnet()
    fids = []
    fid_value = calculate_fid_given_paths(model,
                                          f"fid_images2/base.pt",
                                          "trained_models/mix_len_cr_percent_equal_batch_3_13_prop_100_hdsz_2048_batch_2000_optim_adam_lr_0.001_wd_0.0_enocoderdrop_0.0_drop_0.2_step_mix_mode_12.pth",
                                          # "random_images.pt",
                                          100,
                                          2048)
    print(fid_value)
    # for i in range(17):
    #     fid_value = calculate_fid_given_paths(model,
    #                                           f"fid_images2/{i}.pt",
    #                                           # "trained_models/mix_len_cr_percent_equal_batch_3_13_prop_100_hdsz_2048_batch_2000_optim_adam_lr_0.001_wd_0.0_enocoderdrop_0.0_drop_0.2_step_mix_mode_12.pth",
    #                                           "random_images.pt",
    #                                           100,
    #                                           2048)
    #
    #     print('FID: ', fid_value)
    #     fids.append(fid_value)

    # fig, ax = plt.subplots()
    # ax.plot(fids)
    # plt.savefig("fids.png")
    # with open("fids_random.txt", "w") as file:
    #     for f in fids:


import os
import sys
import numpy as np
from math import floor, ceil, sqrt
from scipy.signal import convolve

class PDB(object):
    def __init__(self, pdb_file):
        self.pdb_file = pdb_file
        self.coords = []
        self.info = []

        if not os.path.exists(self.pdb_file):
            print("PDB> File not found: %s"%self.pdb_file)
            sys.exit(1)

        '''
        Load Coordinate sections of a PDB file as of PDB version 3.30

        COLUMNS        DATA  TYPE    FIELD        DEFINITION
        -------------------------------------------------------------------------------------
         1 -  6        Record name   "ATOM  "
         7 - 11        Integer       serial       Atom  serial number.
        13 - 16        Atom          name         Atom name.
        17             Character     altLoc       Alternate location indicator.
        18 - 20        Residue name  resName      Residue name.
        22             Character     chainID      Chain identifier.
        23 - 26        Integer       resSeq       Residue sequence number.
        27             AChar         iCode        Code for insertion of residues.
        31 - 38        Real(8.3)     x            Orthogonal coordinates for X in Angstroms.
        39 - 46        Real(8.3)     y            Orthogonal coordinates for Y in Angstroms.
        47 - 54        Real(8.3)     z            Orthogonal coordinates for Z in Angstroms.
        55 - 60        Real(6.2)     occupancy    Occupancy.
        61 - 66        Real(6.2)     tempFactor   Temperature  factor.
        77 - 78        LString(2)    element      Element symbol, right-justified.
        79 - 80        LString(2)    charge       Charge  on the atom.
        '''
        self.CA_idx = []
        self.BB_idx = []
        c = 0
        with open(self.pdb_file, 'r') as pdb:
            for line in pdb:
                line_type = line[0:6].strip()
                if line_type in ["ATOM", "HETATM"]:
                    try:
                        at_num = int(line[6:11].strip())
                        at_name = line[12:16].strip()
                        res_name = line[17:20]
                        chain_id = line[21]
                        res_num = int(line[22:26].strip())
                        x = float(line[30:38])
                        y = float(line[38:46])
                        z = float(line[46:54])
                        element_symbol = line[76:78].strip()
                    except Exception:
                        # print("PDB> Error while reading line, skipping it: \n     > %s"%line)
                        pass

                    self.info.append([at_num, at_name, res_name, chain_id, res_num, element_symbol, line_type])
                    self.coords.append([x, y, z])
                    if at_name == "CA":
                        self.CA_idx.append(c)
                    if at_name in ["C", "CA", "N", "O"]:
                        self.BB_idx.append(c)
                    c += 1

        self.coords = np.array(self.coords)
        self.CA_idx = tuple(self.CA_idx)
        self.n_atoms = len(self.coords)
        self.n_CA = len(self.CA_idx)

        # Extent of structure
        self.minx = np.amin(self.coords[:, 0])
        self.maxx = np.amax(self.coords[:, 0])
        self.miny = np.amin(self.coords[:, 1])
        self.maxy = np.amax(self.coords[:, 1])
        self.minz = np.amin(self.coords[:, 2])
        self.maxz = np.amax(self.coords[:, 2])

    def write_pdb(self, outname):
        occupancy = 1.0
        tempFactor = 0.0
        fout = open(outname, "w")
        for i in range(self.n_atoms):
            # Check for atom type: if 4 letters long, start in column 13 otherwise no space.
            # Else, start in column 14. Differs from format convention v3.30 but always see it this way in files.
            if len(self.info[i][1]) == 4:
                line_model = "%-6s%5i %-4s %3s%2s%4s    %8.3f%8.3f%8.3f%6.2f%6.2f          %-2s"
            else:
                line_model = "%-6s%5i  %-3s %3s%2s%4s    %8.3f%8.3f%8.3f%6.2f%6.2f          %-2s"
            line = line_model%(self.info[i][6], self.info[i][0], self.info[i][1], self.info[i][2], self.info[i][3], self.info[i][4], self.coords[i][0], self.coords[i][1], self.coords[i][2], occupancy, tempFactor, self.info[i][5])

            fout.write(line+"\n")
        fout.close()

    def rgyr(self):
        rgyr = np.sqrt(np.sum(np.sum((self.coords - self.center())**2, axis=1) / self.coords.shape[0]))
        return rgyr

    ####################
    # PDB manipulation #
    ####################
    def get_coords(self):
        return self.coords

    def set_coords(self, coords):
        self.coords = coords.copy()

    def rotate_atoms(self, rot_mat):
        self.coords = np.dot(self.coords, rot_mat)

    def translate_atoms(self, trans_vec):
        self.coords += np.array(trans_vec)

    def get_rmsd_with(self, pdb):
        dist_sq = np.square(self.coords - pdb.coords)
        return np.sqrt(np.sum(dist_sq, axis=(0,1)) / dist_sq.shape[0])

    def get_rmsdCA_with(self, pdb):
        if not len(self.CA_idx):
            print("PDB> No alpha carbons detected; returning all-atom RMSD instead.")
            return self.get_rmsd_with(pdb)
        dist_sq = np.square(self.coords[self.CA_idx, :] - pdb.coords[pdb.CA_idx, :])
        return np.sqrt(np.sum(dist_sq, axis=(0,1)) / dist_sq.shape[0])

    ###################
    # Density related #
    ###################
    # Get density grid
    def structure_to_density(self, resolution, voxelsp, isovalue=0.0, pad=0, outname=""):
        # Interpolate pdb to grid + box size + origin coordinates
        grid_prot, pxb, pyb, pzb, minx, miny, minz = self.interpolate_to_grid_massweighted(voxelsp, pad=pad)

        # Margin for convoluted box
        margin = 2 + pad

        # Sigma
        # > Divide by pi*sqrt(2): it is a factor that affects the voxelsp of the gaussian.
        #   Same factor as in chimera. Their explanation in documentation:
        #       Fourier transform (FT) of the distribution fall to 1/e of its maximum value at wavenumber 1/resolution
        # > Divide by voxel spacing to make sure kernel relates to structure space
        # > Truncate values to 3 sigma
        sig = resolution / (np.pi*sqrt(2)) / voxelsp
        r = int(ceil(3.0 * sig))  # truncate at 3 sigma

        # Make Gaussian kernel
        z,y,x = np.ogrid[-r:r+1, -r:r+1, -r:r+1]
        h = np.exp( -(x*x + y*y + z*z) / (2. * sig**2) )
        kernel = h / h.sum()

        # Make density map at given resolution from structure grid
        grid_prot = np.reshape(grid_prot, (pxb, pyb, pzb), order='F')
        grid_dens = convolve(grid_prot, kernel).astype(np.float32)

        # Origin and size of convoluted map
        dxi = minx - (r + margin) * voxelsp;
        dyi = miny - (r + margin) * voxelsp;
        dzi = minz - (r + margin) * voxelsp;

        # Bring map in range [0,1] + apply isovalue threshold
        grid_dens = np.divide(grid_dens, np.amax(grid_dens))
        grid_dens[np.where(grid_dens < isovalue)] = 0

        if outname != "":
            extension = os.path.splitext(outname)[-1].lower()
            if extension in [".sit", ".situs"]:
                dxb, dyb, dzb = grid_dens.shape
                f_out = open(outname,"w")
                f_out.write("%f %f %f %f %i %i %i\n\n"%(voxelsp, dxi, dyi, dzi, dxb, dyb, dzb))
                voxi = 0
                for z in range(dzb):
                    for y in range(dyb):
                        for x in range(dxb):
                            if (voxi+1)%10 == 0 :
                                f_out.write("\n")
                            f_out.write("   %6.6f   "%grid_dens[x][y][z])
                            voxi += 1
                f_out.close()
            else:
                # assuming mrc/map
                import mrcfile
                with mrcfile.new(outname, overwrite=True) as mrc:
                    mrc.set_data(grid_dens.transpose(2,1,0).astype(np.float32))
                    xb, yb, zb = grid_dens.shape

                    mrc.mode = 2

                    mrc.header.mx = xb
                    mrc.header.my = yb
                    mrc.header.mz = zb

                    mrc.header.nxstart = 0
                    mrc.header.nystart = 0
                    mrc.header.nzstart = 0
                    mrc.header.origin.x = dxi
                    mrc.header.origin.y = dyi
                    mrc.header.origin.z = dzi

                    mrc.header.cella.x = xb * voxelsp
                    mrc.header.cella.y = yb * voxelsp
                    mrc.header.cella.z = zb * voxelsp

                    mrc.header.mapc = 1
                    mrc.header.mapr = 2
                    mrc.header.maps = 3

        return grid_dens.astype(np.float32), dxi, dyi, dzi

    #################
    # Interpolation #
    #################

    def interpolate_to_grid_massweighted(self, voxelsp, pad=0):
        # Convert 3D idx to 1D
        def idz(kx, ky, k, j, i):
            return kx * ky * k + kx * j + i

        mass_dict = {"H" : 1.00797, "BE": 9.01218, "C" : 12.011, "N" : 14.0067, "O" : 15.9994, "F": 18.998403, "S" : 32.06, "P" : 30.97376,
                     "MG": 24.305, "CL": 35.453, "K": 39.0983, "CA": 40.078, "MN": 54.9380, "FE": 55.847, "NI": 58.70, "CU": 63.546, "ZN": 65.38, "SE": 78.96}

        # To build the density map, we need the XYZ coordinates and mass of for each atom from PDB
        info_for_dens = []
        element_list = [x[-2].upper() for x in self.info]
        for xyz, elem in zip(self.coords, element_list):
            x, y, z = xyz
            if elem in mass_dict.keys():
                mass = mass_dict[elem]
            else:
                print("PDB> (dens) Element %s not in dict. Using mass of carbon."%elem)
                mass = mass_dict["C"]  # default value
            info_for_dens.append([x, y, z, mass])
        info_for_dens = np.array(info_for_dens)

        # Extent of structure
        minx = np.amin(info_for_dens[:, 0])
        maxx = np.amax(info_for_dens[:, 0])
        miny = np.amin(info_for_dens[:, 1])
        maxy = np.amax(info_for_dens[:, 1])
        minz = np.amin(info_for_dens[:, 2])
        maxz = np.amax(info_for_dens[:, 2])

        # Lattice into register with origin

        minx = voxelsp * floor(minx / voxelsp)
        maxx = voxelsp * ceil(maxx / voxelsp)
        miny = voxelsp * floor(miny / voxelsp)
        maxy = voxelsp * ceil(maxy / voxelsp)
        minz = voxelsp * floor(minz / voxelsp)
        maxz = voxelsp * ceil(maxz / voxelsp)

        # Size of box in each direction to contain structure
        margin = 2 + pad
        pxb = ceil((maxx - minx) / voxelsp) + 2 * margin + 1
        pyb = ceil((maxy - miny) / voxelsp) + 2 * margin + 1
        pzb = ceil((maxz - minz) / voxelsp) + 2 * margin + 1

        nvox_prot = pxb * pyb * pzb
        grid_prot = np.zeros((nvox_prot, 1))

        # Trilinear interpolation of structure to box
        for i in range(len(info_for_dens)):
            # position within grid
            gx = margin + (info_for_dens[i][0] - minx) / voxelsp
            gy = margin + (info_for_dens[i][1] - miny) / voxelsp
            gz = margin + (info_for_dens[i][2] - minz) / voxelsp

            # XYZ coords surrounding atom position
            x0 = floor(gx)
            y0 = floor(gy)
            z0 = floor(gz)
            x1 = x0 + 1;
            y1 = y0 + 1;
            z1 = z0 + 1;

            # interpolate
            a = x1 - gx
            b = y1 - gy
            c = z1 - gz
            grid_prot[idz(pxb, pyb, z0, y0, x0)] += info_for_dens[i][3] * a * b * c
            grid_prot[idz(pxb, pyb, z1, y0, x0)] += info_for_dens[i][3] * a * b * (1 - c)
            grid_prot[idz(pxb, pyb, z0, y1, x0)] += info_for_dens[i][3] * a * (1 - b) * c
            grid_prot[idz(pxb, pyb, z0, y0, x1)] += info_for_dens[i][3] * (1 - a) * b * c
            grid_prot[idz(pxb, pyb, z1, y1, x0)] += info_for_dens[i][3] * a * (1 - b) * (1 - c)
            grid_prot[idz(pxb, pyb, z0, y1, x1)] += info_for_dens[i][3] * (1 - a) * (1 - b) * c
            grid_prot[idz(pxb, pyb, z1, y0, x1)] += info_for_dens[i][3] * (1 - a) * b * (1 - c)
            grid_prot[idz(pxb, pyb, z1, y1, x1)] += info_for_dens[i][3] * (1 - a) * (1 - b) * (1 - c)

        grid_prot = np.divide(grid_prot, np.amax(grid_prot))

        return grid_prot, pxb, pyb, pzb, minx, miny, minz


"""
Prototype of a "Shakemap-Lite" system for retreiving expected GMPE
values from an earthquake
"""
import os
import subprocess
import h5py
import numpy as np
import pandas as pd
from openquake.hazardlib.geo import Point, PlanarSurface, MultiSurface, Mesh
from openquake.hazardlib.gsim import get_available_gsims
from openquake.hazardlib.imt import PGA, PGV, SA, from_string
from openquake.hazardlib.contexts import (ContextMaker, get_distances,
                                          SitesContext, DistancesContext)
from openquake.hazardlib.source.rupture import ParametricProbabilisticRupture
from openquake.hazardlib.site import Site, SiteCollection
from openquake.hazardlib.scalerel.wc1994 import WC1994
from openquake.hazardlib import const
import synthetic_rupture_generator as srg

GSIM_SET = get_available_gsims()

def create_planar_surface(top_centroid, strike, dip, area, aspect):
    """
    Given a central location, create a simple planar rupture
    :param top_centroid:
        Centroid of trace of the rupture, as instance of :class:
            openquake.hazardlib.geo.point.Point
    :param float strike:
        Strike of rupture(Degrees)
    :param float dip:
        Dip of rupture (degrees)
    :param float area:
        Area of rupture (km^2)
    :param float aspect:
        Aspect ratio of rupture

    :returns: Rupture as an instance of the :class:
        openquake.hazardlib.geo.surface.planar.PlanarSurface
    """
    rad_dip = dip * pi / 180.
    width = sqrt(area / aspect)
    length = aspect * width
    # Get end points by moving the top_centroid along strike
    top_right = top_centroid.point_at(length / 2., 0., strike)
    top_left = top_centroid.point_at(length / 2.,
                                     0.,
                                     (strike + 180.) % 360.)
    # Along surface width
    surface_width = width * cos(rad_dip)
    vertical_depth = width * sin(rad_dip)
    dip_direction = (strike + 90.) % 360.

    bottom_right = top_right.point_at(surface_width,
                                      vertical_depth,
                                      dip_direction)
    bottom_left = top_left.point_at(surface_width,
                                    vertical_depth,
                                    dip_direction)

    # Create the rupture
    return PlanarSurface(strike, dip, top_left, top_right,
                         bottom_right, bottom_left)

def vs30_to_z1pt0_cy14(vs30, japan=False):
    """
    Returns the estimate depth to the 1.0 km/s velocity layer based on Vs30
    from Chiou & Youngs (2014) California model

    :param numpy.ndarray vs30:
        Input Vs30 values in m/s
    :param bool japan:
        If true returns the Japan model, otherwise the California model
    :returns:
        Z1.0 in m
    """
    if japan:
        c1 = 412. ** 2.
        c2 = 1360.0 ** 2.
        return np.exp((-5.23 / 2.0) *
                      np.log((np.power(vs30,2.) + c1) / (c2 + c1)))
    else:
        c1 = 571 ** 4.
        c2 = 1360.0 ** 4.
        return np.exp((-7.15 / 4.0) * np.log((vs30 ** 4. + c1) / (c2 + c1)))

def vs30_to_z2pt5_cb14(vs30, japan=False):
    """
    Converts vs30 to depth to 2.5 km/s interface using model proposed by
    Campbell & Bozorgnia (2014)

    :param vs30:
        Vs30 values (numpy array or float)

    :param bool japan:
        Use Japan formula (True) or California formula (False)

    :returns:
        Z2.5 in km
    """
    if japan:
        return np.exp(5.359 - 1.102 * np.log(vs30))
    else:
        return np.exp(7.089 - 1.144 * np.log(vs30))

def get_vs30_sites_from_bbox(bbox, isep="\t"):
    """
    Returns a basic site dictionary from a bbox [llon, ulon, llat, ulat]
    """
    filepath=os.path.dirname(__file__)
    site_data_path = os.path.join(filepath, "global_vs30.grd")
    tempfile = "tempfile.grd"
    # Call grdcut
    cutstring = "/".join([str(loc) for loc in bbox])
    subprocess.run(["gmt", "grdcut", site_data_path,
                    "-G{:s}".format(tempfile),
                    "-R{:s}".format(cutstring)])
    # Call grd2xyz
    tempxyz = os.path.join(filepath,"tempfile.xyz")
    subprocess.run(["gmt", "grd2xyz", tempfile, "-sa", ">", tempxyz])
    # Use pandas to read in the xyzdata
    site_data = pd.read_csv(tempxyz, sep=isep)
    sites = {"lon": site_data.iloc[:, 0].values,
             "lat": site_data.iloc[:, 1].values,
             "vs30": site_data.iloc[:, 2].values}
    os.remove(tempfile)
    os.remove(tempxyz)
    return sites

class Event(object):
    """
    Shakemap event object - requires a minimum of an ID, longitude, latitude,
    depth and magnitude. Other attributes con be input to control the rupture
    orientation and mechanism.

    Can input a rupture geometry directly
    """
    def __init__(self, i_d, lon, lat, hypo_depth, mag, strike=0.0, dip=90.,
                 rake=0.0, aspect=1.0, msr=srg.Stafford2014(), usd=0.0,
                 lsd=1000.0, rupture=None):

        self.id = i_d
        self.lon = lon
        self.lat = lat
        self.depth = hypo_depth
        self.mag = mag
        self.strike = strike
        self.dip = dip
        self.rake = rake
        self.mechanism = []
            #({"strike": strike, "dip": dip, "rake": rake}, 1.)]
        self.aspect = aspect
        self.msr = msr
        if rupture:
            assert isinstance(rupture, ParametricProbabilisticRupture)
            self.rupture = rupture
        else:
            self.rupture = None
        self.usd = usd
        self.lsd = lsd
        self.generator = srg.FiniteRuptureSampler(msr, self.usd, self.lsd)

    def __repr__(self):
        return "{:s}|{:.5f}|{:.5f}|{:.5f}|{:.2f}".format(self.id,
                                                         self.lon,
                                                         self.lat,
                                                         self.depth,
                                                         self.mag)
    def get_rupture(self):
        """
        If a rupture is provided then it is returned, otherwise it will
        use a synthetic rupture generator
        """
        if self.rupture:
           return rupture
        else:
            if len(self.mechanism) >= 1:
                # Define a set of rupture mechanisms with weight
                mechanisms = []
                for mech, weight in self.mechanism:
                    npd = NodalPlane(self.mechanism["strike"],
                                     self.mechanism["dip"],
                                     self.mechanism["rake"])
                    mechanisms.append((weight, npd))
                mechanisms = PMF(mechanisms)
            else:
                if self.rake:
                    if self.rake >= 45.0 and self.rake <= 135.0:
                        mechanisms = "R"
                    elif self.rake >= -135.0 and self.rake <= -45.:
                        mechanisms = "N"
                    else:
                        mechanisms = "SS"
                else:
                    mechanisms = "U"
            # Build planar rupture
            planar_surface = self.generator.sample_ruptures(
                Point(self.lon, self.lat, self.depth),
                      self.mag, nsamples=1, mechanisms=mechanisms,
                      dimensions=None)[0]
            return ParametricProbabilisticRupture(
                self.mag, self.rake, None,
                Point(self.lon, self.lat, self.depth),
                planar_surface, 1.0, None)

class ShakemapLite(object):
    """
    Class to implement a lightweight OpenQuake-based shakemap
    """
    def __init__(self, database, gmpes, imts):
        """
        Shakemap results are stored to an hdf5 databse
        :param str database:
            Path to database
        :param list gmpes:
            List of gmpes (as strings)
        :param list imts:
            List of IMTs (as strings)
        """

        self.gmpes = [GSIM_SET[gmpe]() for gmpe in gmpes]
        self.imts = [from_string(imt) for imt in imts]
        self.db_file = database
        self.context = ContextMaker(self.gmpes)
        # Determine the site attributes
        self.site_attribs = []
        for gmpe in self.gmpes:
            for site_attrib in gmpe.REQUIRES_SITES_PARAMETERS:
                if not site_attrib in self.site_attribs:
                    self.site_attribs.append(site_attrib)

    def __call__(self, event, sites, vs30measured=False, backarc=False):
        """
        Execute a shakemap
        """
        # Build the contexts
        sctx, dctx = self.build_sites_distances_contexts(event, sites,
                                                         vs30measured,
                                                         backarc)
        # Connect to database
        fle = h5py.File(self.db_file)
        grp = fle.create_group(event.id)
        # Store distance calculations to database
        dists_grp = grp.create_group("Distances")
        for key, distance in list(dctx.__dict__.items()):
            dset = dists_grp.create_dataset(key, distance.shape, dtype="f")
            dset[:] = distance
        # Get the rupture
        rctx = event.get_rupture()
        # Run the calculations
        for gmpe in self.gmpes:
            gmpe_grp = grp.create_group(gmpe.__class__.__name__)
            for imt in self.imts:
                imt_grp = gmpe_grp.create_group(str(imt))
                # Get ground motion values
                mean, [sigma] = gmpe.get_mean_and_stddevs(
                    sctx, rctx, dctx, imt, [const.StdDev.TOTAL])
                mean_dset = imt_grp.create_dataset("median",
                                                   mean.shape,
                                                   dtype="f")
                mean_dset[:] = np.exp(mean)
                sigma_dset = imt_grp.create_dataset("sigma",
                                                    sigma.shape,
                                                    dtype="f")
                sigma_dset[:] = sigma
        fle.close()

    def reset(self):
        """
        Remove the database file then delete the object itself
        """
        if os.path.exists(self.db_file):
            os.remove(self.db_file)

    def __getitem__(self, i):
        """
        Returns a given data set on a path
        """
        if not os.path.exists(self.db_file):
            raise AttributeError("Database %s does not exist" % self.db_file)
        # Connect to the database
        fle = h5py.File(self.db_file)
        try:
            data = fle[i][:]
            return data
        except:
            # Data not found
            fle.close()
            raise AttributeError("%s not found in database" % i)

    def build_sites_distances_contexts(self, event, sites, vs30measured=False,
                                       backarc=False):
        """
        Builds the contexts from the event and sites
        """
        # Get distances
        mesh = Mesh(sites["lon"], sites["lat"])
        mshape = mesh.lons.shape
        dctx = DistancesContext()
        for param in self.context.REQUIRES_DISTANCES:
            setattr(dctx, param, get_distances(event.get_rupture(),
                                               mesh, param))
        # Get sites context
        sctx = SitesContext()
        for key in self.site_attribs:
            if key.startswith("z1pt0") and not "z1pt0" in sites:
                setattr(sctx, "z1pt0",
                        vs30_to_z1pt0_cy14(sites["vs30"], japan=False))

            elif key.startswith("z2pt5") and not "z5pt5" in sites:
                setattr(sctx, "z2pt5",
                        vs30_to_z2pt5_cb14(sites["vs30"], japan=False))

            elif key.startswith("vs30measured") and not "vs30measured" in sites:
                if vs30measured:
                    setattr(sctx, "vs30measured", np.ones(mshape, dtype=bool))
                else:
                    setattr(sctx, "vs30measured", np.zeros(mshape, dtype=bool))
            elif key.startswith("vs30measured") and not "vs30measured" in sites:
                if vs30measured:
                    setattr(sctx, "vs30measured", np.ones(mshape, dtype=bool))
                else:
                    setattr(sctx, "vs30measured", np.zeros(mshape, dtype=bool))
            elif key.startswith("backarc") and not "backarc" in sites:
                if vs30measured:
                    setattr(sctx, "backarc", np.ones(mshape, dtype=bool))
                else:
                    setattr(sctx, "backarc", np.zeros(mshape, dtype=bool))
            else:
                setattr(sctx, key, sites[key])
        return sctx, dctx

#
#    def build_site_collection(self, sites, vs30measured=False, backarc=False):
#        """
#        Sites should be a dictionary of site parameter
#        {lon, lat, vs30, vs30measured, z1pt0, z2pt5, backarc}
#        """
#        for key in ["lon", "lat", "vs30"]:
#            # check in dict
#            assert key in sites
#        nsites = len(sites["lon"])
#        if not "z1pt0" in sites or sites["z1pt0"] is None:
#            sites["z1pt0"] = vs30_to_z1pt0_cy14(sites["vs30"], japan=False)
#
#        if not "z2pt5" in sites or sites["z2pt5"] is None:
#            sites["z2pt5"] = vs30_to_z2pt5_cb14(sites["vs30"], japan=False)
#
#        if not "vs30measured" in sites or sites["vs30measured"] is None:
#            if vs30measured:
#                sites["vs30measured"] = np.ones(nsites, dtype=bool)
#            else:
#                sites["vs30measured"] = np.zeros(nsites, dtype=bool)
#
#        if not "backarc" in sites or sites["backarc"] is None:
#            sites["backarc"] = backarc
#        site_col = []
#        for i in range(len(sites["lon"])):
#            site = Site(Point(sites["lon"][i], sites["lat"][i], 0.0),
#                        sites["vs30"][i], sites["vs30measured"][i],
#                        sites["z1pt0"][i], sites["z2pt5"][i],
#                        sites["backarc"][i])
#            site_col.append(site)
#        return SiteCollection(site_col)


#!/usr/bin/env python3

"""
This module provides classes used to
  * define a phonon mode
  * apply symmetry op to a phonon mode
"""

import numpy as np
from numpy.linalg import norm

from .interface_vasp import Poscar
from .cluster import Cluster
from .coord_utils import ReadPBC2Cart, match_p1p0
from .util.mathtool import relativePosition
import logging
logger = logging.getLogger(__name__)

debug_level = 10

class PhononMode():
    """

    """

    def __init__(self, prim, clus, u, clus_dominant):
        """
        a proper cluster decorated with displacement u
        """
        self.prim = prim
        self.cluster = clus
        self.u = u
        self.dominant_cluster = clus_dominant
        # print(self.cluster)
        # print(self.dominant_cluster)

    def operated_by(self, isym):
        return PhononMode(self.prim, Cluster.from_ijkl(self.cluster.operated_by(isym), self.prim),
                          np.dot(self.u, self.prim.spacegroup[isym].rot.T),
                          Cluster.from_ijkl(self.dominant_cluster.operated_by(isym), self.prim))

    def __repr__(self):
#        outs = [""]
#        outs.append(self.frac_coords.__str__())
#        outs.append("ijk= "+ self.ijk.__repr__()+ " l="+ self.l.__repr__()+"]")
#        outs.append(str(self.ijkl))
#        return " ".join(outs)
        return str(self.cluster)+ repr(self.u)

    @classmethod
    def from_file(cls, prim, p1f, disp_cut, natom_dominant=2, p0f=None, dominant_specie=None, dominant_dr=True):
        """

        :param p0: read an ideal structure,
         p1: disturbed structure
         disp_cut: retain only atoms with displacement larger than this cutoff
         dominant_specie: which atoms may be the representative
        :paramdominant_dr: Select dominant pairs by largest change in distance
        :return: a phonon mode
        """
        if p0f is None:
            x0 = Poscar.from_file(p1f).structure.map_to_prim(prim)
        else:
            assert isinstance(p0f, str), TypeError("require p0f filename")
            x0 = Poscar.from_file(p0f).structure
        scmat = prim.get_scmat(x0.lattice.matrix)

        if debug_level > 10:
            print("supercell clusters generated")

        u = ReadPBC2Cart(p1f, x0.frac_coords)
        # center of mass
        u -= np.average(u, axis=0)
        unorm = norm(u, axis=1)
        # print(unorm)
        # print(u)
        uidx = np.where(unorm > disp_cut)[0]
        # print(uidx)
        if len(uidx) < natom_dominant:
            print("ERROR: found only %d displaced atoms" % (len(uidx)))
            exit(-1)

        clus_xyz = x0.frac_coords[uidx].copy()
        # move cluster atoms close to one atom considering PBC
        imax = np.argmax(unorm)
        pc0 = x0.frac_coords[imax:imax+1].copy()
        for i in range(len(uidx)):
            # print("debug", i, ui, clus_xyz[i], clus_xyz[i:i+1], pc0)
            clus_xyz[i] = match_p1p0(clus_xyz[i:i+1], pc0)[0]
            # print(clus_xyz[i], x0.frac_coords[ui])
        clus_xyz = np.dot(clus_xyz, scmat)

        # now find out the dominant cluster
        if dominant_dr and natom_dominant == 2:
            max_dr = -1.
            sel = [True for _ in uidx]
            if isinstance(dominant_specie, list):
                ele = x0.elements
                sel = [ele[i] in dominant_specie for i in uidx]
#            print(sel, ele)
            for i1, i in enumerate(uidx):
                if not sel[i1]:
                    continue
                for i2, j in enumerate(uidx):
                    if (not sel[i2]) or (i1 >= i2):
                        continue
                    r0 = prim.lattice.norm(clus_xyz[i2]-clus_xyz[i1])
                    r1 = np.linalg.norm(prim.lattice.get_cartesian_coords(clus_xyz[i2]-clus_xyz[i1]) + u[j]-u[i])
                    dr = abs(r0-r1)
                    if dr > max_dr:
                        max_dr = dr
#                        print(i, j, max_dr)
                        atom_dominant = np.array([i1, i2])
        else:
            if dominant_specie is None:
                udominant = unorm[uidx]
            else:
                ele = x0.elements
                assert isinstance(dominant_specie, list)
                # print("debug species=", dominant_specie, ele, "for finding dimer |u|=", [unorm[i] if ele[i] in dominant_specie else -1. for i in uidx])
                udominant = [unorm[i] if ele[i] in dominant_specie else -1. for i in uidx]
            atom_dominant = np.argsort(udominant)[-natom_dominant:]
#        print('dominant=', atom_dominant)
        for i in uidx[atom_dominant]:
            print("atom %4d element %3s   disp %8.3f" % (i, x0[i].specie, unorm[i]), x0.frac_coords[i])

        return cls(prim, Cluster.from_coords(clus_xyz, prim), u[uidx],
                   Cluster.from_coords(clus_xyz[atom_dominant], prim))

    def apply_to(self, dest, dest_clus, inv_sc_mat):
        """
        Apply phonon mode to a supercell
        :param dest: destination supercell
        :param dest_clus: dominant cluster by which to identify or match the phonon
        :param inv_sc_mat: inverse sc_mat belonging to the supercell
        :return:
        """
        prim = self.prim
        syms = prim.spacegroup
        for isym, sym in enumerate(syms):
            s_pl = self.operated_by(isym)
            is_eq, tr = dest_clus.equivalent_by_lattranslation(s_pl.dominant_cluster.ijkls)
            if is_eq:
                # print(isym, tr, s_pl.dominant_cluster.ijkls, dest_dimer.ijkls)
                # found a sym op to move polaron to desired position
                break
        assert is_eq, ValueError("cannot apply phonon mode")
        pi = relativePosition(dest_clus.ijkls, tr)
        tr_ijk = np.array(dest_clus.ijkls[0]) - np.array(s_pl.dominant_cluster.ijkls[pi[0]])
        s_pl.cluster.move_to_prim(dijk=tr_ijk[:3])
        dest_polaron = [dest.frac2ijkl(np.dot(s_pl.cluster.frac_coords[i], inv_sc_mat))[3] for i in
                        range(s_pl.cluster.order)]

        out_sc = dest.copy()
        for i, isite in enumerate(dest_polaron):
            out_sc[isite].move_by(s_pl.u[i])
        return out_sc