import argparse
import os
import shutil
import numpy as np
import pandas as pd
import numpy as np
import scipy.spatial.distance as ssd
import gc
import h5py
import itertools
import scipy.stats as stats
from scipy.io import mmread
from scipy import sparse, io
from scipy.sparse import csr_matrix, triu, coo_matrix
import numpy as np

class HiCNorm(object):
    def __init__(self, matrix):
        self.bias = None
        self.matrix = matrix
        self.norm_matrix = None
        self._bias_var = None

    def iterative_correction(self, max_iter=400):
        #mat = self.matrix  * 1.0  ## use float to avoid int overflow
        print (np.amax(self.matrix))
        print (np.amin(self.matrix))
        mat_lil = sparse.lil_matrix(self.matrix)
        # remove low count rows for normalization
        row_sum = np.asarray(np.sum(mat_lil, axis=1)).ravel()
        low_count = np.quantile(row_sum, 0.15)  # 0.15 gives best correlation between KR/SP/VC
        map_resolution = np.quantile(row_sum, 0.20)  # map resolution
        print (map_resolution)
        mask_row = row_sum < low_count

        #mask_row = np.asarray(mask_row.ravel())
        #mask_row = [count for count, value in enumerate(mask_row) if value]
        #print (mask_row)
        #mat_lil = sparse.lil_matrix(mat)
        mat_lil[mask_row, :] = 0
        mat_lil[:, mask_row] = 0 

        
        self.bias = np.ones(mat_lil.shape[0])
        print (self.bias.dtype)
        self._bias_var = []
        x, y = np.nonzero(mat_lil)
        mat_sum = np.sum(mat_lil)   # force the sum of matrix to stay the same, otherwise ICE may not converge as fast

        for i in range(max_iter):
            #bias =sparse.csr_matrix.sum(mat, axis=1)
            bias = np.asarray(np.sum(mat_lil, axis=1)).ravel()
            bias_mean = np.mean(bias[bias > 0])
            bias_var = np.var(bias[bias > 0])
            self._bias_var.append(bias_var)
            bias = bias / bias_mean
            bias[bias == 0] = 1

            #print (np.asarray(bias)[:,:].ravel())
            #print (bias[x])
            #print (bias[x]*bias[y])
            
            mat = sparse.csr_matrix(mat_lil)

            mat[x, y] = mat[x, y] / (bias[x]*bias[y])
            #print (mat)
            new_sum = np.sum(mat)
            mat = mat * (mat_sum / new_sum)
            mat_lil = sparse.lil_matrix(mat)

            # update total bias
            self.bias = self.bias * bias * np.sqrt(new_sum / mat_sum)
            if bias_var < .001:
                break
            else:
                pass   
        
        print (self._bias_var[-1])
        self.norm_matrix = csr_matrix(self.matrix,  dtype=np.float64)
        self.norm_matrix[x, y] = self.norm_matrix[x, y] / (self.bias[x] * self.bias[y])
        
        
    def vanilla_coverage(self):
        self.iterative_correction(max_iter=1)


#have to ensure that the input matrix is a square matrix
def lib_3d_norm_old(arr):
    s = sparse.csr_matrix.get_shape(arr)[0]
    arr = sparse.dia_matrix(arr, dtype=np.float32)
    lib_norm_hic = sparse.dia_matrix((s, s), dtype=np.float32)

    for K in list(range(1, s)):
        diagonal_array = sparse.dia_matrix.diagonal(arr, k=K)
        diagonal_mean = np.nanmean(diagonal_array)
        if diagonal_mean == 0:
            pass
        else:
            sparse.dia_matrix.setdiag(lib_norm_hic, (diagonal_array/diagonal_mean), k=K)
    return sparse.csr_matrix(lib_norm_hic)

def lib_3d_norm(arr, resolution):
    s = arr.shape
    if s[0] != s[1]:
        raise ValueError('network must be a square matrix')
    else:
        s = s[0]
    arr = sparse.dia_matrix(arr, dtype=np.float32)
    lib_norm_hic = sparse.dia_matrix((s, s), dtype=np.float32)
    L = int(250000/ resolution)
    for K in list(range(1, s)):
        diagonal_array = sparse.dia_matrix.diagonal(arr, k=K)
        counter = 0
        while diagonal_array.shape[0] < L:
            counter = counter + 1
            if (K + counter) < s:
                diagonal_array = np.concatenate((diagonal_array, sparse.dia_matrix.diagonal(arr, k=(K+counter))), axis=0)
            else:
                pass 
            diagonal_array = np.concatenate((diagonal_array, sparse.dia_matrix.diagonal(arr, k=(K-counter))), axis=0)
        diagonal_mean = np.nanmean(diagonal_array[0:L])
        if diagonal_mean == 0:
            pass
        else:
            sparse.dia_matrix.setdiag(lib_norm_hic, (diagonal_array/diagonal_mean), k=K)
    return sparse.csr_matrix(lib_norm_hic)

def rank_norm(arr, resolution):
    s = arr.shape
    if s[0] != s[1]:
        raise ValueError('network must be a square matrix')
    else:
        s = s[0]
    arr = sparse.dia_matrix(arr, dtype=np.float32)
    rank_norm_hic = sparse.dia_matrix((s, s), dtype=np.float16)
    L = int(250000/ resolution)
    L = int(197200/ resolution)
    for K in list(range(1, s)):
        diagonal_array = sparse.dia_matrix.diagonal(arr, k=K)
        counter = 0
        while diagonal_array.shape[0] < L:
            counter = counter + 1
            if (K + counter) < s:
                diagonal_array = np.concatenate((diagonal_array, sparse.dia_matrix.diagonal(arr, k=(K+counter))), axis=0)
            else:
                pass 
            diagonal_array = np.concatenate((diagonal_array, sparse.dia_matrix.diagonal(arr, k=(K-counter))), axis=0)

        rank_abs = lambda x: stats.rankdata(np.abs(x))
        normalized_diag = np.apply_along_axis(rank_abs, 0, diagonal_array[0:L])
        sparse.dia_matrix.setdiag(rank_norm_hic, normalized_diag, k=K)                                    
    return sparse.csr_matrix(rank_norm_hic)