import cartopy.crs as projections
import numpy as np
import matplotlib.pyplot as plt
from datetime import timedelta
import xarray as xr

from typhon.plots import worldmap
from typhon.collocations import Collocator, expand, collapse
from typhon.files import FileSet, NetCDF4
from typhon.collocations import Collocations

"""
Explanation: Find collocations with typhon
End of explanation
"""


# Create the data
primary = xr.Dataset(
    coords={
        "lat": (('along_track'), 30.*np.sin(np.linspace(-3.14, 3.14, 24))+20),
        "lon": (('along_track'), np.linspace(0, 90, 24)),
        "time": (('along_track'), np.arange("2018-01-01", "2018-01-02", dtype="datetime64[h]")),
    },
    data_vars={
        "Temperature": (("along_track"), np.random.normal(290, 5, (24))),
    }
)
secondary = xr.Dataset(
    coords={
        "lat": (('along_track'), 30.*np.sin(np.linspace(-3.14, 3.14, 24)+1.)+20),
        "lon": (('along_track'), np.linspace(0, 90, 24)),
        "time": (('along_track'), np.arange("2018-01-01", "2018-01-02", dtype="datetime64[h]")),
    },
    data_vars={
        "Temperature": (("along_track"), np.random.normal(290, 5, (24))),
    }
)

# Plot the data
fig = plt.figure(figsize=(10, 10))
wmap = worldmap(primary["lat"], primary["lon"], s=24, bg=True)
worldmap(secondary["lat"], secondary["lon"], s=24, ax=wmap.axes,)

"""
Explanation: Collocations between two data arrays
Let's try out the simplest case: You have two xarray datasets with
temporal-spatial data and you want to find collocations between them.
At first, we create two example xarray datasets with faked measurements. Let's
assume, these data arrays represent measurements from two different instruments
(e.g. on satellites). Each measurement has a time attribute indicating when
it was taken and a geo-location (latitude and longitude) indicating where
this happened. Note that the lat and lon variables must share their first
dimension with the time coordinate.
End of explanation
"""


collocator = Collocator(name='primary_secondary_collocator')
collocations = collocator.collocate(
    primary=('primary', primary),
    secondary=('secondary', secondary),
    max_distance=600, # collocation radius in km
)
print(f'Found collocations are {collocations["Collocations/distance"].values} km apart')

collocations

"""
Explanation: Now, let’s find all measurements of primary that have a maximum distance of 300 kilometers to the measurements of secondary:
End of explanation
"""


def collocations_wmap(collocations):
    fig = plt.figure(figsize=(10, 10))

    # Plot the collocations
    wmap = worldmap(
        collocations['primary/lat'],
        collocations['primary/lon'],
        facecolor="r", s=128, marker='x', bg=True
    )
    worldmap(
        collocations['secondary/lat'],
        collocations['secondary/lon'],
        facecolor="r", s=128, marker='x', bg=True, ax=wmap.axes
    )

    # Plot all points:
    worldmap(primary["lat"], primary["lon"], s=24, ax=wmap.axes,)
    worldmap(secondary["lat"], secondary["lon"], s=24, ax=wmap.axes,)
    wmap.axes.set(ylim=[-15, 55], xlim=[-10, 100])

collocations_wmap(collocations)

"""
Explanation: The obtained collocations dataset contains variables of 3 groups: primary, secondary and Collocations.
The first two correspond to the variables of the two respective input datasets and contain only the matched 
data points. The Collocations group adds some new variables containing information about the collocations, e.g.
the temporal and spatial distances. Additional information can be found in the typhon documentation
Let’s mark the collocations with red crosses on the map:
End of explanation
"""


collocations = collocator.collocate(
    primary=('primary', primary),
    secondary=('secondary', secondary),
    max_distance=300, # collocation radius in km
    max_interval=timedelta(hours=1), # temporal collocation interval as timedelta
)
print(
    f'Found collocations are {collocations["Collocations/distance"].values} km apart in space '
    f'and {collocations["Collocations/interval"].values} hours apart in time.'
)
collocations_wmap(collocations)

"""
Explanation: We can also add a temporal filter that filters out all points which difference in time is bigger than a time interval. We are doing this by using max_interval. Note that our testdata is sampled very sparsely in time.
End of explanation
"""


expand(collocations)

"""
Explanation: As mentioned in :func:collocate, the collocations are returned in compact format, e.g. an efficient way to store the collocated data. When several data points in the secondary group collocate with a single observation of the primary group, it is not obvious how this should be handled. The compact format accounts for this by introducing the Collocations/pairs variable, which contains the respective indices of the collocated datapoints. This might not be the most practical solution 
In practice, the two functions expand and collapse offer two convenient ways to handle this. 
Applying expand to the collocations will repeat data points for cases where one datapoint matches with several data points of the other dataset.
End of explanation
"""


collapse(collocations)

"""
Explanation: Applying collapse to the collocations will calculate some generic statistics (mean, std, count) over the datapoints that match with a single data point of the other dataset.
End of explanation
"""


fh = NetCDF4()
fh.write(secondary, 'testdata/secondary/2018/01/01/000000-235959.nc')

# Create the filesets objects and point them to the input files
a_fileset = FileSet(
    name="primary",
    path="testdata/primary/{year}/{month}/{day}/"
         "{hour}{minute}{second}-{end_hour}{end_minute}{end_second}.nc",
#     handler=handlers.NetCDF4,
)
b_fileset = FileSet(
    name="secondary",
    path="testdata/secondary/{year}/{month}/{day}/"
         "{hour}{minute}{second}-{end_hour}{end_minute}{end_second}.nc",
#     handler=handlers.NetCDF4,
)

"""
Explanation: Purely temporal collocations are not implemented yet and attempts will raise a NotImplementedError.
Find collocations between two filesets
Normally, one has the data stored in a set of many files. typhon provides an object to handle those filesets (see the typhon doc). It is very simple to find collocations between them.
First, we need to create FileSet objects and let them know where to find their files:
End of explanation
"""


# Create the output dataset:
ab_collocations = Collocations(
    name="ab_collocations",
    path="testdata/ab_collocations/{year}/{month}/{day}/"
       "{hour}{minute}{second}-{end_hour}{end_minute}{end_second}.nc",   
)
ab_collocations.search(
    [a_fileset, b_fileset], start="2018", end="2018-01-02",
    max_interval=timedelta(hours=1), max_distance=300, 
)

fh.read('testdata/primary/2018/01/01/000000-235959.nc')

"""
Explanation: Now, we can search for collocations between a_dataset and b_dataset and store them to ab_collocations.
End of explanation
"""


