# -*- coding: utf-8 -*-
# coding: utf-8
import numpy as np
import pyproj
from GrASP import Params_GrASP


def Make_Stadis_array(Station_loc_data):
    """ 
    Make Distance matrix

    Parameter:
    ----------
    Station_loc_data: Station locations

    Return:
    ----------
    Stadis_array: Distance matrix (unit: km)
    """

    Stadis_array = np.zeros((np.shape(Station_loc_data)[0], np.shape(Station_loc_data)[0]))
    transformer = pyproj.Transformer.from_crs(Params_GrASP.EPSG_base, Params_GrASP.EPSG_proj, always_xy = True) 
    locx, locy = transformer.transform(Station_loc_data[:,2], Station_loc_data[:,1])
    for i in range(np.shape(Station_loc_data)[0]):
        Dis = 0.001*np.sqrt((locx - locx[i])**2 + (locy - locy[i])**2) ### km unit
        Stadis_array[i,:] = Dis

    return Stadis_array


##### k nearest neighobor #####
def kNN_graph(Stadis_array, k_neareast = 10):
    """
    Make kNN graph (Recommended !!!)

    Parameters:
    ----------
    Stadis_array: Distance matrix
    k_neareast: k of kNN (default: k=10)

    Return:
    ----------
    Output_array: Adjacency matrix
    """

    Output_array = np.zeros(np.shape(Stadis_array))
    Graph_mode_name = 'K_'+str(k_neareast)
    Stadis_array_nearest = Stadis_array.copy()
    for i in range(np.shape(Stadis_array_nearest)[0]):
        sort_list = sorted (Stadis_array_nearest[i,:])
        k_thres = sort_list[k_neareast]
        Nearest_list = Stadis_array_nearest[i,:].copy()
        Nearest_list[Nearest_list > k_thres] = 0
        Nearest_list[Nearest_list > 0] = 1
        Output_array[i,:] += Nearest_list
        Output_array[:,i] += Nearest_list
    Output_array[Output_array > 0] = 1
    print(Graph_mode_name)

    return Output_array


##### epsilon nearest neighobor #####
def ENN_graph(Stadis_array, epsilon = 100):
    """ 
    Make epsilon-NN graph

    Parameters:
    ----------
    Stadis_array: Distance matrix
    epsilon: radius (km) (default: epsilon=100)

    Return:
    ----------
    Output_array: Adjacency matrix
    """

    Output_array = Stadis_array.copy()
    Graph_mode_name = 'epsilon_' + str(epsilon)
    Output_array[Stadis_array > epsilon] = 0
    Output_array[Output_array > 0] = 1
    print(Graph_mode_name)
    
    return Output_array


##### mutual k nearest neighobor #####
def Mutual_kNN_graph(Stadis_array, k_neareast = 10):
    """ 
    Make mutual kNN graph (properties intermediate between kNN and ENN)

    Parameters:
    ----------
    Stadis_array: Distance matrix
    k_neareast: k of kNN (default: k=10)

    Return:
    ----------
    Output_array: Adjacency matrix
    """

    Output_array = np.ones(np.shape(Stadis_array))
    Graph_mode_name = 'MK_'+str(k_neareast)
    Stadis_array_nearest = Stadis_array.copy()
    for i in range(np.shape(Stadis_array_nearest)[0]):
        sort_list = sorted (Stadis_array_nearest[i,:])
        k_thres = sort_list[k_neareast]
        Nearest_list = Stadis_array_nearest[i,:].copy()
        Nearest_list[Nearest_list > k_thres] = 0
        Nearest_list[Nearest_list > 0] = 1
        Output_array[i,:] *= Nearest_list
        Output_array[:,i] *= Nearest_list
    Output_array[Output_array > 0] = 1
    print(Graph_mode_name)

    return Output_array