# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use('TkAgg')
import matplotlib.pyplot as plt
import math
from ctypes import cdll, c_float, c_double, Structure, ARRAY, POINTER, c_int32, c_uint32, c_size_t, c_void_p, cast
from sys import platform
from bng import bng
import pyproj
import ipdb
from array import array
%matplotlib inline
%config InlineBackend.figure_format='retina'

def helmert(x, y, z=0):
    """ Example implementation of Helmert transform """
    tX = -446.448  
    tY = 125.157  
    tZ = -542.060  

    rX = -0.1502  
    rY = -0.2470  
    rZ = -0.8421  

    s = 20.4894 * math.pow(10, -6)
    # For multiple x and y
    # A_vector = np.matrix([[x, y, z], [x, y, z]]).T
    A_vector = np.vstack(np.array([x, y, z]))
    t_matrix = np.vstack(np.array([tX, tY, tZ]))
    conversion = np.matrix([[1 + s, -rZ, rY], [rZ, 1 + s, -rX], [-rY, rX, 1 + s]])
    return t_matrix + (conversion * A_vector)

helmert(-2.018304, 54.589097)

"""
Explanation: The Helmert transform
$
\begin{bmatrix}
    x\
    y\
    z\
\end{bmatrix}^B
=
\begin{bmatrix}
    t_X\
    t_Y\
    t_Z\
\end{bmatrix}
+
\begin{bmatrix}
    1+s  & -r_Z & r_Y\
    r_Z  & 1+s  & -r_X\
    -r_Y & r_X  & 1+s\
\end{bmatrix}
\cdot
\begin{bmatrix}
    x\
    y\
    z\
\end{bmatrix}^A
$
Using an external Rust library to speed up lon, lat to BNG conversion
End of explanation
"""


if platform == "darwin":
    ext = "dylib"
else:
    ext = "so"
    
lib = cdll.LoadLibrary('target/release/liblonlat_bng.' + ext)

"""
Explanation: Setting up the Rust library. See here for more
Ensure you've built your Rust library using cargo build --release, or the next step will fail.
The boilerplate below can easily be hidden in a wrapper function – it's just here to demonstrate how to call into a shared Rust lib using FFI.
End of explanation
"""


class _FFIArray(Structure):
    """ Convert sequence of floats to a C-compatible void array """
    _fields_ = [("data", c_void_p),
                ("len", c_size_t)]

    @classmethod
    def from_param(cls, seq):
        """  Allow implicit conversions from a sequence of 64-bit floats."""
        return seq if isinstance(seq, cls) else cls(seq)

    def __init__(self, seq, data_type = c_double):
        """
        Convert sequence of values into array, then ctypes Structure

        Rather than checking types (bad), we just try to blam seq
        into a ctypes object using from_buffer. If that doesn't work,
        we try successively more conservative approaches:
        numpy array -> array.array -> read-only buffer -> CPython iterable
        """
        if isinstance(seq, float):
            seq = array('d', [seq])
        try:
            len(seq)
        except TypeError:
             # we've got an iterator or a generator, so consume it
            seq = array('d', seq)
        array_type = data_type * len(seq)
        try:
            raw_seq = array_type.from_buffer(seq.astype(np.float64))
        except (TypeError, AttributeError):
            try:
                raw_seq = array_type.from_buffer_copy(seq.astype(np.float64))
            except (TypeError, AttributeError):
                # it's a list or a tuple
                raw_seq = array_type.from_buffer(array('d', seq))
        self.data = cast(raw_seq, c_void_p)
        self.len = len(seq)
        

class _Result_Tuple(Structure):
    """ Container for returned FFI data """
    _fields_ = [("e", _FFIArray),
                ("n", _FFIArray)]


def _void_array_to_list(restuple, _func, _args):
    """ Convert the FFI result to Python data structures """
    eastings = POINTER(c_double * restuple.e.len).from_buffer_copy(restuple.e)[0]
    northings = POINTER(c_double * restuple.n.len).from_buffer_copy(restuple.n)[0]
    res_list = [list(eastings), list(northings)]
    drop_array(restuple.e, restuple.n)
    return res_list

"""
Explanation: Define the ctypes structures for lon, lat --> BNG conversion
End of explanation
"""


# Multi-threaded FFI functions
convert_bng = lib.convert_to_bng_threaded
convert_bng.argtypes = (_FFIArray, _FFIArray)
convert_bng.restype = _Result_Tuple
convert_bng.errcheck = _void_array_to_list
convert_bng.__doc__ = """
    Multi-threaded lon, lat --> BNG conversion
    Returns a list of two lists containing Easting and Northing floats,
    respectively
    Uses the Helmert transform
    """

convert_lonlat = lib.convert_to_lonlat_threaded
convert_lonlat.argtypes = (_FFIArray, _FFIArray)
convert_lonlat.restype = _Result_Tuple
convert_lonlat.errcheck = _void_array_to_list
convert_lonlat.__doc__ =  """
    Multi-threaded BNG --> lon, lat conversion
    Returns a list of two lists containing Longitude and Latitude floats,
    respectively
    Uses the Helmert transform
    """

convert_to_osgb36 = lib.convert_to_osgb36_threaded
convert_to_osgb36.argtypes = (_FFIArray, _FFIArray)
convert_to_osgb36.restype = _Result_Tuple
convert_to_osgb36.errcheck = _void_array_to_list
convert_to_osgb36.__doc__ = """
    Multi-threaded lon, lat --> OSGB36 conversion, using OSTN02 data
    Returns a list of two lists containing Easting and Northing floats,
    respectively
    """

convert_osgb36_to_lonlat = lib.convert_osgb36_to_ll_threaded
convert_osgb36_to_lonlat.argtypes = (_FFIArray, _FFIArray)
convert_osgb36_to_lonlat.restype = _Result_Tuple
convert_osgb36_to_lonlat.errcheck = _void_array_to_list
convert_osgb36_to_lonlat.__doc__ = """
    Multi-threaded OSGB36 --> Lon, Lat conversion, using OSTN02 data
    Returns a list of two lists containing Easting and Northing floats,
    respectively
    """

convert_etrs89_to_lonlat = lib.convert_etrs89_to_ll_threaded
convert_etrs89_to_lonlat.argtypes = (_FFIArray, _FFIArray)
convert_etrs89_to_lonlat.restype = _Result_Tuple
convert_etrs89_to_lonlat.errcheck = _void_array_to_list
convert_etrs89_to_lonlat.__doc__ = """
    Multi-threaded ETRS89 Eastings and Northings --> OSGB36 conversion, using OSTN02 data
    Returns a list of two lists containing Easting and Northing floats,
    respectively
    """

convert_etrs89_to_osgb36 = lib.convert_etrs89_to_osgb36_threaded
convert_etrs89_to_osgb36.argtypes = (_FFIArray, _FFIArray)
convert_etrs89_to_osgb36.restype = _Result_Tuple
convert_etrs89_to_osgb36.errcheck = _void_array_to_list
convert_etrs89_to_osgb36.__doc__ = """
    Multi-threaded OSGB36 Eastings and Northings --> ETRS89 Eastings and Northings conversion,
    using OSTN02 data
    Returns a list of two lists containing Easting and Northing floats,
    respectively
    """

convert_osgb36_to_etrs89 = lib.convert_osgb36_to_etrs89_threaded
convert_osgb36_to_etrs89.argtypes = (_FFIArray, _FFIArray)
convert_osgb36_to_etrs89.restype = _Result_Tuple
convert_osgb36_to_etrs89.errcheck = _void_array_to_list
convert_osgb36_to_etrs89.__doc__ = """
    Multi-threaded ETRS89 Eastings and Northings --> Lon, Lat conversion,
    Returns a list of two lists containing Longitude and Latitude floats,
    respectively
    """

convert_epsg3857_to_wgs84 = lib.convert_epsg3857_to_wgs84_threaded
convert_epsg3857_to_wgs84.argtypes = (_FFIArray, _FFIArray)
convert_epsg3857_to_wgs84.restype = _Result_Tuple
convert_epsg3857_to_wgs84.errcheck = _void_array_to_list
convert_epsg3857_to_wgs84.__doc__ = """
    Convert Google Web Mercator (EPSG3857) coordinates to WGS84
    Latitude and Longitude
    Returns a list of two lists containing latitudes and longitudes,
    respectively
    """

# Free FFI-allocated memory
drop_array = lib.drop_float_array
drop_array.argtypes = (_FFIArray, _FFIArray)
drop_array.restype = None

"""
Explanation: Define ctypes input and return parameters
End of explanation
"""


# UK bounding box
N = 55.811741
E = 1.768960
S = 49.871159
W = -6.379880

bng = pyproj.Proj(init='epsg:27700')
bng_ostn02 = pyproj.Proj("+init=EPSG:27700 +nadgrids=OSTN02_NTv2.gsb")
wgs84 = pyproj.Proj(init='epsg:4326')

num_coords = 1000000
lon_ls = list(np.random.uniform(W, E, [num_coords]))
lat_ls = list(np.random.uniform(S, N, [num_coords]))

"""
Explanation: Simple test of average conversion speed, Python version
Test: 1MM random points within the UK
End of explanation
"""


%%timeit -r50
[bng(lat, lon) for lat, lon in zip(lat_ls, lon_ls)]

"""
Explanation: Pure Python
End of explanation
"""


%%timeit -r50
pyproj.transform(wgs84, bng, lon_ls, lat_ls)

"""
Explanation: Pyproj
End of explanation
"""


%%timeit -r50
convert_bng(lon_ls, lat_ls)

"""
Explanation: Multithreaded Rust
End of explanation
"""


%run remote_bench.py

# make graph text look good - only run this if you know what you're doing
from matplotlib import rc
rc('font', **{'family':'sans-serif',
    'sans-serif':['Helvetica'],
    'monospace': ['Inconsolata'],
    'serif': ['Helvetica']})
rc('text', **{'usetex': True})
rc('text', **{'latex.preamble': '\\usepackage{sfmath}'})

df = pd.read_csv("benches/benchmarks.csv",
    dtype={
        'crossbeam': np.float64,
        'crossbeam_error': np.float64,
        'threads': np.float64,
        'rayon': np.float64,
        'rayon_error': np.float64,
        'weight': np.float64,
        'num_points': np.float64,
        'cores': np.float64,
    }    
)

df.head(20)

"""
Explanation: Pyproj is now only 1.15x (~15%) faster than Multithreaded Rust, which is 9x faster than pure Python
Benchmark using cProfile
End of explanation
"""


plt.clf()
fig = plt.figure(figsize=(12, 8))

# Crossbeam
# filter by cores and method
df_cb = df[(df.cores == 2) & (df.method == "crossbeam")]
ax1 = fig.add_subplot(121, facecolor='#d3d3d3')
threads = plt.scatter(df_cb['ivariable'], df_cb['time'],
    color='#CC00CC',
    edgecolor='#333333',
    marker='o',
    lw=1,
    s=80,
    alpha=1.0,
    zorder=2,
)
ax1.set_ylim(0, 40000000)
plt.errorbar(df_cb['ivariable'], df_cb['time'], yerr=df_cb['error'],
    color='#ff96ca',
    linestyle="None",
    lw=2.,
    alpha=0.75,
    zorder=1     
)
# filter by cores and method (8 cores)
df_cb_ = df[(df.cores == 8) & (df.method == "crossbeam")]
threads_ = plt.scatter(df_cb_['ivariable'], df_cb_['time'],
    color='#008080',
    edgecolor='#333333',
    marker='o',
    lw=1,
    s=80,
    alpha=1.0,
    zorder=2,
)
plt.errorbar(df_cb_['ivariable'].values, df_cb_['time'].values, yerr=df_cb_['error'].values,
    linestyle="None",
    color='#cc8400',
    lw=2.,
    alpha=0.75,
    zorder=1     
)

plt.ylabel("Time (ns)", fontsize=14)
plt.xlabel("Num. threads", fontsize=14)
# highlight fastest runs
rd = plt.axhline(df_cb.time.min(), color='red')
rd_ = plt.axhline(df_cb_.time.min(), color='red')
# Legend
leg = plt.legend(
    (threads, threads_, rd),
    ('2 cores, 1.8GHz Core i7', '8 cores, 3.4GHz Core i7', 'Fastest run'),
    loc='upper right',
    scatterpoints=1,
    fontsize=9)
leg.get_frame().set_alpha(0.5)
ax1.spines['top'].set_visible(False)
ax1.spines['right'].set_visible(False)

# x ticks
plt.tick_params(
    axis='x',
    which='both',
    top=False,
    bottom='on',
    labelbottom='on')

# y ticks
plt.tick_params(
    axis='y',
    which='both',
    left='on',
    right=False,
    labelbottom=False)


plt.title("Crossbeam, 100k Points", fontsize=14)
# ax1.tick_params(axis='both', which='major', labelsize=12)
plt.tight_layout()

# Rayon
# filter by cores and method
df_ry = df[(df.cores == 2) & (df.method == "rayon")]
ax2 = fig.add_subplot(122, facecolor='#d3d3d3')
rayon = plt.scatter(df_ry['ivariable'], df_ry['time'],
    color='#CC00CC',
    edgecolor='#000000',
    marker='o',
    lw=1,
    s=60,
    alpha=1.0,
    zorder=2,
)
ax2.set_ylim(0, 40000000)
# pandas bug, so use .values
plt.errorbar(df_ry['ivariable'].values, df_ry['time'].values, yerr=df_ry['error'].values,
    linestyle="None",
    color='#ff96ca',
    lw=2.,
    alpha=0.75,
    zorder=1 
)

# filter by cores and method (8 cores)
df_ry_ = df[(df.cores == 8) & (df.method == "rayon")]
rayon_ = plt.scatter(df_ry_['ivariable'], df_ry_['time'],
    color='#008080',
    edgecolor='#333333',
    marker='o',
    lw=1,
    s=80,
    alpha=1.0,
    zorder=2,
)
plt.errorbar(df_ry_['ivariable'].values, df_ry_['time'].values, yerr=df_ry_['error'].values,
    linestyle="None",
    color='#cc8400',
    lw=2.,
    alpha=0.75,
    zorder=1     
)

plt.xlabel("Weight", fontsize=14)
# ax2.tick_params(axis='both', which='major', labelsize=12)
ax2.set_yticklabels([])
# highlight fastest runs
ry = plt.axhline(df_ry.time.min(), color='red')
ry_ = plt.axhline(df_ry_.time.min(), color='red')
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(True)
ax2.spines['left'].set_visible(False)
plt.title("Rayon, 100k Points", fontsize=14)

leg = plt.legend(
    (rayon, rayon_, ry),
    ('2 cores, 1.8GHz Core i7', '8 cores, 3.4GHz Core i7', 'Fastest run'),
    loc='upper right',
    scatterpoints=1,
    fontsize=9)
leg.get_frame().set_alpha(0.5)

# x ticks
plt.tick_params(
    axis='x',
    which='both',
    top=False,
    bottom='on',
    labelbottom='on')

#  y ticks
plt.tick_params(
    axis='y',
    which='both',
    left=False,
    right=False,
    top=False,
    labelbottom=False)

# output
plt.tight_layout()
plt.savefig(
    'crossbeam_v_rayon.png',
    format="png", bbox_inches='tight',
    alpha=True, transparent=True, dpi=100)
plt.show()

"""
Explanation: Comparing and Optimising Crossbeam and Rayon
End of explanation
"""


