import os
import csv
import numpy as np
from eqep.interpolation.rbf_pgvinterpolator import RbfPGVInterpolator


class EarthQuake:
    """Stores the data of an earthquake

    It has a grid of coordinates with a certain (interpolated)
    peak ground velocity for every cell.

    This class is usually used in combination with a `PGVInterpolator`:

    >>> earthquake = EarthQuake((0, 10), (0, 10), (10, 10))
    >>> earthquake.interpolate_pgvs(eqdata, RbfPGVInterpolator())

    Attributes
    ----------
    longs, lats : array of float
        store the longitudes and latitudes of the grid
    pgvs : array of float
        stores the peak ground velocities corresponding to the grid cells
    """

    def __init__(self, long_range, lat_range, steps,
                 pgvs=None):
        """Creates a new `EarthQuake` for the specified area

        Parameters
        ----------
        long_range : tuple of float
            the longitude range: (min, max)
        lat_range : tuple of float
            the latitude range: (min, max)
        steps : tuple of int
            the number of steps in these ranges
        pgvs : array of floats, optional
            the peak ground velocities corresponding to the grid defined with
            the other parameters
        """

        # initialize the longs and lats
        self.longs, self.lats = np.meshgrid(np.linspace(long_range[0],
                                                        long_range[1],
                                                        steps[0]),
                                            np.linspace(lat_range[0],
                                                        lat_range[1],
                                                        steps[1]))

        if pgvs is None or pgvs.shape == self.longs.shape:
            self.pgvs = pgvs
        else:
            raise AttributeError('`pgvs` must equal the shape of the '
                                      'long- and lat-grid')

    def interpolate_pgvs(self, data, interpolator=RbfPGVInterpolator()):
        """Interpolates the pgvs for the whole earthquake grid

        Parameters
        ----------
        data : EQData
            the data to interpolate to a full earthquake
        interpolator : PGVInterpolator
            the algorithm to use for interpolation
        """

        self.pgvs = interpolator.interpolate(data, self.longs, self.lats)

    def save_csv(self, filename):
        """Saves this object as a CSV-file

        The data is persisted in the following format: ::

            # comment
            long, lat, pgv
            long, lat, pgv
            ...

        Parameters
        ----------
        filename : str
            the name of the file ending with ".csv" (may also be a path)
        """

        if self.pgvs is None:
            raise ValueError('`pgvs` must not be "None".')

        # check if path to filename already exists, if not create it
        if '/' in filename and not os.path.exists(os.path.dirname(filename)):
            os.makedirs(os.path.dirname(filename))

        # write header comment and the values to csv file
        with open(filename, 'w') as file:
            file.write('# longitude, latitude, peak ground velocity\n\n')
            csv_writer = csv.writer(file, lineterminator='\n')

            for long, lat, pgv in zip(self.longs.flatten(), self.lats.flatten(),
                                      self.pgvs.flatten()):
                csv_writer.writerow([long, lat, pgv])
