#!/usr/bin/env python
# coding: utf-8

# In[1]:


import sys
sys.path.append("..")
import gc
import os
import psutil
import pathlib

import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import pearsonr

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint
from torch.optim.lr_scheduler import ReduceLROnPlateau

#import selene_sdk
#from selene_sdk.samplers.dataloader import SamplerDataLoader

#from selene_utils2 import *
from G3D_selene_utils import *
from G3D_mat_module import Decoder, motif

import subprocess
import warnings
warnings.filterwarnings('ignore')
batch_size=1
TF_num=128
AP_num=10
ORCA_PATH = str(pathlib.Path.cwd().parent)
print(ORCA_PATH)

def array32(*args, **kwargs):
    return np.array(*args, dtype=np.float32, **kwargs)
    
def array16(*args, **kwargs):
    return np.array(*args, dtype=np.float16, **kwargs)

def get_gpu_memory_usage():
    try:
        nvidia_smi_output = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,nounits,noheader"], encoding="utf-8")

        gpu_memory = []
        for line in nvidia_smi_output.strip().split("\n"):
            used_memory, total_memory = line.split(",")
            gpu_memory.append({"used": int(used_memory.strip()), "total": int(total_memory.strip())})
        for i, gpu in enumerate(gpu_memory):
            print(f"GPU {i}: Used Memory = {gpu['used']/1024:.2f} GB, Total Memory = {gpu['total']/1024:.2f} GB")
        # return gpu_memory
    except subprocess.CalledProcessError as e:
        print("Error executing nvidia-smi:", e)
        return []

def figshow(x, vmax, np=False, save_file="./tmp.pdf", show=True, vmin=0):
    if np:
        plt.imshow(x.squeeze(), vmin=vmin, vmax=vmax)
    else:
        plt.imshow(x.squeeze().cpu().detach().numpy(), vmin=vmin, vmax=vmax)
    plt.colorbar()
    plt.savefig(save_file)
    if show:
        plt.show()
    plt.close()

modelstr = "h1esc_2M"
seed = 1

torch.set_default_dtype(torch.float32)
os.makedirs("./models/", exist_ok=True)
os.makedirs("./png/", exist_ok=True)

below_20k_ratio = 0.583
below_2m_ratio = 0.803
scale_factor = (1 - below_20k_ratio) / (below_2m_ratio - below_20k_ratio)
print(scale_factor)#比例因子，原数据未剔除>2Mb的数据，这个因子用于修正
# The sumnorm.mcool file contains 20k~2m and >2m contacts. The scale factor need to be multiplied to get sumnorm of 20k~2m contacts

t = Genomic2DFeatures(
    ["../resources/H1_HESC.filter20k.sumnorm.mcool::/resolutions/4000"],
    ["r4000"],
    (500, 500),
    cg=False,
)
print(t)
#sampler是随机获取的？
sampler = RandomPositionsSamplerHiC(
    reference_sequence=MemmapGenome(
        input_path="../resources/Homo_sapiens.GRCh38.dna.primary_assembly.fa",
        memmapfile="../resources/Homo_sapiens.GRCh38.dna.primary_assembly.fa.mmap",
        blacklist_regions="hg38",
    ),
    target=t,
    features=["r4000"],
    test_holdout=["chr9", "chr10"],
    validation_holdout=["chr8"],
    sequence_length=2000000,
    position_resolution=4000,
    random_shift=1000,
    random_strand=True,
    cross_chromosome=False,
)
print(sampler.target)
sampler.mode = "train"
dataloader = SamplerDataLoader(sampler, num_workers=0, batch_size=1, seed=seed)
print(dataloader)
#似乎没有权重给我们用捏

net0 = nn.DataParallel(motif(TF_num, AP_num))
denet = nn.DataParallel(Decoder(TF_num, AP_num, 1, 0))

net0.cuda()
denet.cuda()

net0.train()
denet.train()

params = (
    [p for p in net0.parameters() if p.requires_grad]
    + [p for p in denet.parameters() if p.requires_grad]
)
optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.98)
optimizer_path = "./models/model_h1esc_2M.optimizer"
"""
if os.path.isfile(optimizer_path):
    print("optimizer exist, loading")
    optimizer_state_dict = torch.load(optimizer_path)
    optimizer.load_state_dict(optimizer_state_dict)
scheduler = ReduceLROnPlateau(optimizer, mode="max", factor=0.9, patience=10, threshold=0)
"""
i = 0

sequence = torch.zeros((batch_size, 2000000, 4), dtype=torch.float32)
target = torch.zeros((batch_size, 500, 500), dtype=torch.float32)

get_gpu_memory_usage()


# In[2]:
TF_exp=torch.ones(TF_num)
AP_exp=torch.ones(AP_num)

valid_mses = {1: [], 2: []}
valid_corrs = {1: [], 2: []}
past_losses=[]

# In[ ]:

这部分在初始化模型，可以跑通

while True:
    fillind = 0
    optimizer.zero_grad()
    for ss, tt in dataloader:
        if np.isnan(tt.numpy()).mean() > 0.5:
            continue
        sequence[fillind, :] = ss
        target[fillind, :] = tt 
        if fillind == sequence.shape[0] - 1:
            fillind = 0
        else:
            fillind += 1
            continue
            
        encoding = net0(torch.Tensor(sequence.float()).transpose(1, 2).cuda())
        def train_step(level, start, coarse_pred=None):
            target_np = target[
                :, start : start + 250 * level, start : start + 250 * level
            ].numpy()
            target_r = np.nanmean(
                np.nanmean(
                    np.reshape(target_np, (target.shape[0], 500, 1, 500, 1)), axis=4
                ),
                axis=2,
            )

            pred = denet.forward(
                TF_num, TF_exp, AP_exp, encoding,
            )
            target_cuda = torch.Tensor(target).cuda()
            loss = (
                (
                    pred[:, :, :][~torch.isnan(target_cuda)]
                    - target_cuda[~torch.isnan(target_cuda)]
                )
                ** 2
            ).mean()
            past_losses.append(loss.detach().cpu().numpy())
            return loss, pred, target_r

        start = 0
        loss, pred2, target2 = train_step(2, start)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        del encoding

        frequen=100
        if i % frequen == 0:
            print("l2:" + str(np.mean(past_losses[-frequen:])), flush=True)
            data_to_write="第{}个样本loss:{}\n".format(i,loss)
            with open('example.txt', 'a', encoding='utf-8') as file:
                file.write(data_to_write)
            vmax = np.mean([torch.max(pred2).cpu().detach().numpy(), np.max(target2[0, :, :])])
            figshow(pred2[0, :, :], vmax = vmax, save_file="./png/model_" + modelstr + "." + str(i) + ".level2.pred.pdf", show=False)
            figshow(target2[0, :, :], vmax = vmax, np=True, save_file="./png/model_" + modelstr + "." + str(i) + ".level2.label.pdf", show=False)
            
            torch.save(net0.state_dict(), "./models/model_" + modelstr + ".checkpoint")
            torch.save(denet.state_dict(), "./models/model_" + modelstr + ".d1.checkpoint")
            torch.save(optimizer.state_dict(), "./models/model_" + modelstr + ".optimizer")

        del pred2

        if i % 5000 == 0 and i != 0:
            t2 = Genomic2DFeatures(
                ["../resources/H1_HESC.filter20k.sumnorm.mcool::/resolutions/4000"],
                ["r4000"],
                (500, 500),
                cg=False,
            )
            t2.mode = "validate"
            sampler2 = RandomPositionsSamplerHiC(
                reference_sequence=MemmapGenome(
                    input_path="../resources/Homo_sapiens.GRCh38.dna.primary_assembly.fa",
                    memmapfile="../resources/Homo_sapiens.GRCh38.dna.primary_assembly.fa.mmap",
                    blacklist_regions="hg38",
                ),
                target=t2,
                features=["r4000"],
                test_holdout=["chr9", "chr10"],
                validation_holdout=["chr8"],
                sequence_length=2000000,
                position_resolution=4000,
                random_shift=0,
                random_strand=True,
                cross_chromosome=False,
            )

            sampler2.mode = "validate"
            dataloader2 = SamplerDataLoader(sampler2, num_workers=0, batch_size=1)
            net0.eval()
            denet.eval()

            count = 0
            mses = {1: [], 2: []}
            corrs = {1: [], 2: []}
            fillind = 0
            for ss, tt in dataloader2:
                if np.isnan(tt.numpy()).mean() > 0.3:
                    continue
                sequence[fillind, :] = ss
                target[fillind, :] = tt * scale_factor
                if fillind == sequence.shape[0] - 1:
                    fillind = 0
                else:
                    fillind += 1
                    continue

                count += 1
                if count == 100:
                    break
                encoding = net0(torch.Tensor(sequence.float()).transpose(1, 2).cuda())

                def eval_step(level, start, coarse_pred=None, plot=False):
                    target_r = np.nanmean(
                        np.nanmean(
                            np.reshape(
                                target[
                                    :, start : start + 250 * level, start : start + 250 * level
                                ].numpy(),
                                (target.shape[0], 250, level, 250, level),
                            ),
                            axis=4,
                        ),
                        axis=2,
                    )

                    pred = denet.forward(
                        TF_num, TF_exp, AP_exp, encoding,
                    )
                    print(pred.shape)
                    target_cuda = torch.Tensor(target_r).cuda()
                    loss = (
                        (
                            pred[:, 0:, 0:][~torch.isnan(target_cuda)]
                            - target_cuda[~torch.isnan(target_cuda)]
                        )
                        ** 2
                    ).mean()

                    mses[level].append(loss.detach().cpu().numpy())
                    valid_mses[level].append(loss.detach().cpu().numpy())
                    pred_np = (
                        pred[:, 0:, 0:].detach().cpu().numpy().reshape((pred.shape[0], -1))
                    )
                    target_np = target_r.reshape((pred.shape[0], -1))

                    for j in range(pred_np.shape[0]):
                        validinds = ~np.isnan(target_np[j, :])
                        if np.mean(validinds) > 0.3:
                            corrs[level].append(pearsonr(pred_np[j, validinds], target_np[j, validinds])[0])
                            valid_corrs[level].append(pearsonr(pred_np[j, validinds], target_np[j, validinds])[0])
                        else:
                            corrs[level].append(np.nan)
                            valid_corrs[level].append(np.nan)
                    if plot:
                        for ii in range(pred.shape[0]):
                            vmax = np.mean([torch.max(pred[ii, 0, :, :]).cpu().detach().numpy(), np.max(target_r[ii, :, :])])
                            figshow(pred[ii, :, :], vmax = vmax, save_file="./png/model_" + modelstr + ".test" + str(ii) + ".level" + str(level) + ".pred.pdf", show=False)
                            figshow(target_r[ii, :, :], vmax = vmax, np=True, save_file="./png/model_" + modelstr + ".test" + str(ii) + ".level" + str(level) + ".lebel.pdf", show=False)
                    return pred

                start = 0
                pred = eval_step(2, start, plot=count == 99)
                start = start + 62 * 2
                pred = eval_step(1, start, pred[:, 62:187, 62:187], plot=count == 99)

                del encoding
                del pred

            scheduler.step(np.nanmean(corrs[2]))
            print(
                "Cor1 {0}, Cor2 {1}".format(
                    np.nanmean(corrs[1]),
                    np.nanmean(corrs[2]),
                )
            )
            print(
                "MSE1 {0}, MSE2 {1}".format(
                    np.nanmean(mses[1]),
                    np.nanmean(mses[2]),
                )
            )
            net0.train()
            denet.train()
        i += 1


# In[ ]:

这部分是在训练，训练完一轮后第二次前向传播可以过，到loss.backward()就开始报错

RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). 
Saved intermediate values of the graph are freed when you call .backward() or autograd.grad().
Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.