import time
import numpy as np
from numpy import mean
from numpy import std
from numpy import max
from numpy import min
from numpy import empty
from numpy import cov
from numpy import argmax
from numpy import argmin
import numpy.linalg as LA
import h5py
import math
from math import sqrt
import numba
from numba import jit
import joblib
from joblib import Parallel, delayed, load


@jit(nopython=True)
# 计算canonical correlation
def get_corr(m):
    cov_m = np.cov(m)
    cov11 = empty((2, 2), dtype=np.complex128)
    cov11[0] = cov_m[0, 0:2]
    cov11[1] = cov_m[1, 0:2]
    cov12 = empty((2, 2), dtype=np.complex128)
    cov12[0] = cov_m[0, 2:4]
    cov12[1] = cov_m[1, 2:4]
    cov21 = empty((2, 2), dtype=np.complex128)
    cov21[0] = cov_m[2, 0:2]
    cov21[1] = cov_m[3, 0:2]
    cov22 = empty((2, 2), dtype=np.complex128)
    cov22[0] = cov_m[2, 2:4]
    cov22[1] = cov_m[3, 2:4]

    # 计算矩阵的特征值得到canonical correlation：此处采用原始矩阵，避免矩阵-1/2次幂的计算
    matrix = (LA.inv(cov11)) @ cov12 @ (LA.inv(cov22)) @ cov21
    eigen_vals = LA.eig(matrix)[0]
    if np.real(eigen_vals[0]) >= np.real(eigen_vals[1]):
        ei_vector = LA.eig(matrix)[1].T[0]
        eigen_val = np.real(eigen_vals[0])
    else:
        ei_vector = LA.eig(matrix)[1].T[1]
        eigen_val = np.real(eigen_vals[1])

    # 通过特征向量alpha和beta的乘积判断正负
    corr = sqrt(eigen_val)
    alpha = ei_vector.T
    beta = LA.inv(cov22) @ cov21 @ alpha / corr
    if np.real(alpha.T @ beta) > 0:
        return corr
    else:
        return -corr


@jit(nopython=True)
# 在不同的time shift下遍历所有数据，通过get_corr求得不同数据之间的canonical correlation
# output1 ~ output6为储存w_max ~ tau_min的memmap文件（一个文件即一个726*726*467的数组）
def calculate_cc(data1, data2, output1, output2, output3, output4, output5, output6, n1, n2):
    w_max_lst = empty(467)
    w_min_lst = empty(467)
    c_max_lst = empty(467)
    c_min_lst = empty(467)
    tau_max_lst = empty(467)
    tau_min_lst = empty(467)
    index_store = 181

    # 以年为单位的time shift
    for i in range(467):
        if ((i % 12 == 0) or (i % 12 == 2) | (i % 12 == 4) | (i % 12 == 5) | (i % 12 == 7) | (i % 12 == 9) | (i % 12 == 11)):
            x = data1[index_store+31: index_store+31+365]
            index_store += 31
        elif ((i % 12 == 1) | (i % 12 == 3) | (i % 12 == 8) | (i % 12 == 10)):
            x = data1[index_store+30: index_store+30+365]
            index_store += 30
        else:
            x = data1[index_store+28: index_store+28+365]
            index_store += 28

        x11 = x.T[0]
        x12 = x.T[1]
        lst = empty(401)
        for tau in range(-200, 201):
            y = data2[index_store+tau: index_store+365+tau]
            y11 = y.T[0]
            y12 = y.T[1]
            m = empty((4, 365))
            m[0] = x11
            m[1] = x12
            m[2] = y11
            m[3] = y12
            lst[tau+200] = get_corr(m)

        c_max = max(lst)
        c_min = min(lst)
        w_max = (c_max-mean(lst))/std(lst)
        w_min = (c_min-mean(lst))/std(lst)
        tau_max = argmax(lst)  # 此处由于numba的一些限制，输出的tau为0-400，对应-200-200。
        tau_min = argmin(lst)

        w_max_lst[i] = w_max
        w_min_lst[i] = w_min
        c_max_lst[i] = c_max
        c_min_lst[i] = c_min
        tau_max_lst[i] = tau_max
        tau_min_lst[i] = tau_min

    output1[n1][n2] = w_max_lst
    output2[n1][n2] = w_min_lst
    output3[n1][n2] = c_max_lst
    output4[n1][n2] = c_min_lst
    output5[n1][n2] = tau_max_lst
    output6[n1][n2] = tau_min_lst


if __name__ == "__main__":
    start_time = time.time()

    # 创建六个memmap数组文件
    W_max = np.memmap('w_max', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    W_min = np.memmap('w_min', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    C_max = np.memmap('c_max', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    C_min = np.memmap('c_min', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    TAU_max = np.memmap('tau_max', dtype=np.float64,
                        shape=(726, 726, 467), mode='w+')
    TAU_min = np.memmap('tau_min', dtype=np.float64,
                        shape=(726, 726, 467), mode='w+')

    # 使用joblib进行并行计算
    Parallel(n_jobs=4)(delayed(calculate_cc)(np.array([data_850[i], data_1000[i]]).T, np.array(
        [data_850[j], data_1000[j]]).T, W_max, W_min, C_max, C_min, TAU_max, TAU_min, i, j) for i in range(726) for j in range(726))

    # 将6个储存w_max ~ tau_min的memmap数组写入新创建的h5文件
    with h5py.File('results.hdf5', 'w', libver='latest') as f:
        w_max = f.create_dataset('w_max', (726, 726, 467), data=W_max)
        w_min = f.create_dataset('w_min', (726, 726, 467), data=W_min)
        c_max = f.create_dataset('c_max', (726, 726, 467), data=C_max)
        c_min = f.create_dataset('c_min', (726, 726, 467), data=C_min)
        tau_max = f.create_dataset('tau_max', (726, 726, 467), data=TAU_max)
        tau_min = f.create_dataset('tau_min', (726, 726, 467), data=TAU_min)
    f.close()

    end_time = time.time()
    print('total running time (min): ', (end_time-start_time)/60, 'min')

import time
import numpy as np
from numpy import mean
from numpy import std
from numpy import max
from numpy import min
from numpy import empty
from numpy import cov
from numpy import argmax
from numpy import argmin
import numpy.linalg as LA
import h5py
import math
from math import sqrt
import numba
from numba import jit
import joblib
from joblib import Parallel, delayed, load


@jit(nopython=True)
# 计算canonical correlation
def get_corr(m):
    cov_m = np.cov(m)
    cov11 = empty((2, 2), dtype=np.complex128)
    cov11[0] = cov_m[0, 0:2]
    cov11[1] = cov_m[1, 0:2]
    cov12 = empty((2, 2), dtype=np.complex128)
    cov12[0] = cov_m[0, 2:4]
    cov12[1] = cov_m[1, 2:4]
    cov21 = empty((2, 2), dtype=np.complex128)
    cov21[0] = cov_m[2, 0:2]
    cov21[1] = cov_m[3, 0:2]
    cov22 = empty((2, 2), dtype=np.complex128)
    cov22[0] = cov_m[2, 2:4]
    cov22[1] = cov_m[3, 2:4]

    # 计算矩阵的特征值得到canonical correlation：此处采用原始矩阵，避免矩阵-1/2次幂的计算
    matrix = (LA.inv(cov11)) @ cov12 @ (LA.inv(cov22)) @ cov21
    eigen_vals = LA.eig(matrix)[0]
    if np.real(eigen_vals[0]) >= np.real(eigen_vals[1]):
        ei_vector = LA.eig(matrix)[1].T[0]
        eigen_val = np.real(eigen_vals[0])
    else:
        ei_vector = LA.eig(matrix)[1].T[1]
        eigen_val = np.real(eigen_vals[1])

    # 通过特征向量alpha和beta的乘积判断正负
    corr = sqrt(eigen_val)
    alpha = ei_vector.T
    beta = LA.inv(cov22) @ cov21 @ alpha / corr
    if np.real(alpha.T @ beta) > 0:
        return corr
    else:
        return -corr


@jit(nopython=True)
# 在不同的time shift下遍历所有数据，通过get_corr求得不同数据之间的canonical correlation
# output1 ~ output6为储存w_max ~ tau_min的memmap文件（一个文件即一个726*726*467的数组）
def calculate_cc(data1, data2, output1, output2, output3, output4, output5, output6, n1, n2):
    w_max_lst = empty(467)
    w_min_lst = empty(467)
    c_max_lst = empty(467)
    c_min_lst = empty(467)
    tau_max_lst = empty(467)
    tau_min_lst = empty(467)
    index_store = 181

    # 以年为单位的time shift
    for i in range(467):
        if ((i % 12 == 0) or (i % 12 == 2) | (i % 12 == 4) | (i % 12 == 5) | (i % 12 == 7) | (i % 12 == 9) | (i % 12 == 11)):
            x = data1[index_store+31: index_store+31+365]
            index_store += 31
        elif ((i % 12 == 1) | (i % 12 == 3) | (i % 12 == 8) | (i % 12 == 10)):
            x = data1[index_store+30: index_store+30+365]
            index_store += 30
        else:
            x = data1[index_store+28: index_store+28+365]
            index_store += 28

        x11 = x.T[0]
        x12 = x.T[1]
        lst = empty(401)
        for tau in range(-200, 201):
            y = data2[index_store+tau: index_store+365+tau]
            y11 = y.T[0]
            y12 = y.T[1]
            m = empty((4, 365))
            m[0] = x11
            m[1] = x12
            m[2] = y11
            m[3] = y12
            lst[tau+200] = get_corr(m)

        c_max = max(lst)
        c_min = min(lst)
        w_max = (c_max-mean(lst))/std(lst)
        w_min = (c_min-mean(lst))/std(lst)
        tau_max = argmax(lst)  # 此处由于numba的一些限制，输出的tau为0-400，对应-200-200。
        tau_min = argmin(lst)

        w_max_lst[i] = w_max
        w_min_lst[i] = w_min
        c_max_lst[i] = c_max
        c_min_lst[i] = c_min
        tau_max_lst[i] = tau_max
        tau_min_lst[i] = tau_min

    output1[n1][n2] = w_max_lst
    output2[n1][n2] = w_min_lst
    output3[n1][n2] = c_max_lst
    output4[n1][n2] = c_min_lst
    output5[n1][n2] = tau_max_lst
    output6[n1][n2] = tau_min_lst


if __name__ == "__main__":
    start_time = time.time()

    # 创建六个memmap数组文件
    W_max = np.memmap('w_max', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    W_min = np.memmap('w_min', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    C_max = np.memmap('c_max', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    C_min = np.memmap('c_min', dtype=np.float64,
                      shape=(726, 726, 467), mode='w+')
    TAU_max = np.memmap('tau_max', dtype=np.float64,
                        shape=(726, 726, 467), mode='w+')
    TAU_min = np.memmap('tau_min', dtype=np.float64,
                        shape=(726, 726, 467), mode='w+')

    # 使用joblib进行并行计算
    Parallel(n_jobs=4)(delayed(calculate_cc)(np.array([data_850[i], data_1000[i]]).T, np.array(
        [data_850[j], data_1000[j]]).T, W_max, W_min, C_max, C_min, TAU_max, TAU_min, i, j) for i in range(726) for j in range(726))

    # 将6个储存w_max ~ tau_min的memmap数组写入新创建的h5文件
    with h5py.File('results.hdf5', 'w', libver='latest') as f:
        w_max = f.create_dataset('w_max', (726, 726, 467), data=W_max)
        w_min = f.create_dataset('w_min', (726, 726, 467), data=W_min)
        c_max = f.create_dataset('c_max', (726, 726, 467), data=C_max)
        c_min = f.create_dataset('c_min', (726, 726, 467), data=C_min)
        tau_max = f.create_dataset('tau_max', (726, 726, 467), data=TAU_max)
        tau_min = f.create_dataset('tau_min', (726, 726, 467), data=TAU_min)
    f.close()

    end_time = time.time()
    print('total running time (min): ', (end_time-start_time)/60, 'min')

# 计算canonical correlation # 计算矩阵的特征值得到canonical correlation：此处采用原始矩阵，避免矩阵-1/2次幂的计算 # 通过特征向量alpha和beta的乘积判断正负 # 在不同的time shift下遍历所有数据，通过get_corr求得不同数据之间的canonical correlation # output1 ~ output6为储存w_max ~ tau_min的memmap文件（一个文件即一个726*726*467的数组） # 以年为单位的time shift # 此处由于numba的一些限制，输出的tau为0-400，对应-200-200。 # 创建六个memmap数组文件 # 使用joblib进行并行计算 # 将6个储存w_max ~ tau_min的memmap数组写入新创建的h5文件