
from get_data import *
from RF import *
import numpy as np
import torch
from sklearn.metrics.cluster import normalized_mutual_info_score
from sklearn.metrics.cluster import adjusted_mutual_info_score
from sklearn.metrics.cluster import adjusted_rand_score
from sklearn.metrics.cluster import rand_score
from scipy.cluster.hierarchy import DisjointSet
import matplotlib.pyplot as plot
import random
from connected_component import *
import copy
import matplotlib.pyplot as plt
import matplotlib.pyplot as plot



RFMS_simple_iteration       = {'UserKnowledge':100   ,'iris':100   ,'WallRobot':100   ,'WirelessLocalization':100  }
RFMS_simple_bandwidth       = {'UserKnowledge':0.5   ,'iris':0.55  ,'WallRobot':0.2   ,'WirelessLocalization':0.7  }
RFMS_simple_threshold       = {'UserKnowledge':0.9   ,'iris':0.9   ,'WallRobot':0.9   ,'WirelessLocalization':0.9  }
RFMS_simple_step_size       = {'UserKnowledge':0.005 ,'iris':0.005 ,'WallRobot':0.005 ,'WirelessLocalization':0.005}
RFMS_simple_learning_rate   = {'UserKnowledge':0.03  ,'iris':0.03  ,'WallRobot':0.01  ,'WirelessLocalization':0.05 }
RFMS_simple_n_input         = {'UserKnowledge':5     ,'iris':4     ,'WallRobot':4     ,'WirelessLocalization':7    }
RFMS_simple_dim             = {'UserKnowledge':2000  ,'iris':2000  ,'WallRobot':500   ,'WirelessLocalization':750  }




def RFMS_shift(data,dim,bandwidth,iteration,learning_rate,smoothing_parameter,QMC=True,blurring=False):
    n, in_dim = data.shape
    if QMC:
        encoder = QMCRF_Encoder(in_dim,dim,bandwidth)
    else:
        encoder = MCRF_Encoder(in_dim,dim,bandwidth)
    encodings = encoder.encode_x(data)
    
    c = (2*np.pi*(bandwidth**2))**(in_dim/2)
    density_encoding = (1 / (n*c)) * encodings.sum(axis=0)
    
    current_position = copy.deepcopy(data)
    current_encoding = copy.deepcopy(encodings)
    
    for t in range(iteration):
        random_direction = np.random.normal(size=(n, in_dim))
        prob = current_position + (smoothing_parameter*random_direction)
        prob_encoding = encoder.encode_x(prob)
        local_density = encoder.similarity(current_encoding,density_encoding)
        prob_density = encoder.similarity(prob_encoding,density_encoding)
        gradient = random_direction * (((prob_density - local_density).real)/smoothing_parameter)[:, np.newaxis]
        update = learning_rate * (gradient / (local_density[:, np.newaxis]))
        current_position = current_position + update
        current_encoding = encoder.encode_x(current_position)
        if blurring:
            density_encoding = (1 / (n*c)) * current_encoding.sum(axis=0)
    
    return current_position



def run_RFMS(X, y,count,dataset):
	n_input = RFMS_simple_n_input[dataset]
	dim = RFMS_simple_dim[dataset]
	bandwidth = RFMS_simple_bandwidth[dataset]
	step_size = RFMS_simple_step_size[dataset]
	learning_rate = RFMS_simple_learning_rate[dataset]
	threshold = RFMS_simple_threshold[dataset]
	iteration = RFMS_simple_iteration[dataset]
	shifted = RFMS_shift(X,dim,bandwidth,iteration,learning_rate,step_size,True,False)

	final_gram = Gaussian_kernel(shifted,shifted,bandwidth)
	pred = CC(final_gram,threshold)

	return normalized_mutual_info_score(pred,y), adjusted_mutual_info_score(pred,y), adjusted_rand_score(pred,y), rand_score(pred,y)



def run_RFMS_blurring(X, y,count,dataset):
	n_input = RFMS_simple_n_input[dataset]
	dim = RFMS_simple_dim[dataset]
	bandwidth = RFMS_simple_bandwidth[dataset]
	step_size = RFMS_simple_step_size[dataset]
	learning_rate = RFMS_simple_learning_rate[dataset]
	threshold = RFMS_simple_threshold[dataset]
	iteration = RFMS_simple_iteration[dataset]
	shifted = RFMS_shift(X,dim,bandwidth,iteration,learning_rate,step_size,True,True)

	final_gram = Gaussian_kernel(shifted,shifted,bandwidth)
	pred = CC(final_gram,threshold)
	return normalized_mutual_info_score(pred,y), adjusted_mutual_info_score(pred,y), adjusted_rand_score(pred,y), rand_score(pred,y)








