# -*- coding: utf-8 -*-
"""
Created on Wed Mar  8 15:30:19 2023

@author: Songjian Lu

---------- Modified on Sep. 13, 2024
Add two line in point-plot: "green"--linear regression with shift; "red"--linear regression without shift.
Add the choice of no-shift for normalization
"""
import numpy, math, shutil
from numpy.linalg import inv
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt 
from matplotlib.figure import Figure



p_st_using = 'Y'

testShow ='N'

def read_meta_info_sample_ctype(p_file, p_cid_colNo = 0, p_sp_colNo=2, p_type_colNo = 1, p_sp = '\t'):
   L_rets = {}; L_ctype = {}
   ff = open(p_file, 'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln0 = ln.split('\n')
         ln1 = ln0[0].split(p_sp)
         ctype = ln1[p_type_colNo]
         cid = ln1[p_cid_colNo]
         spID = ln1[p_sp_colNo]

         L_rets[cid] = [spID, ctype]
         if not ctype in L_ctype:
            L_ctype[ctype] = 0
         L_ctype[ctype] += 1

   ff.close()
   
   if (p_st_using.upper())[0] == 'Y':
      print('--- Counts of different cell types ---')
      print('  cell Type name\tcell number')
   for tp in L_ctype:
      print('  ',tp, '\t\t', L_ctype[tp])

         
   return L_rets

def check_scRNA(p_file,p_st = 0):
   #p_st = 0:org count, 1:log count
   L_rets = {}; L_min = 10000.0; L_max = -1000.0
   L_rets2 = {}
   L_colNo = 0; L_rowNo = 0
   L_title = []
   ff = open(p_file,'r')
   LineNo = 0
   for ln in ff:
      LineNo += 1
   ff.close()
   if (p_st_using.upper())[0] == 'Y':
      print('---- Total row#:', LineNo)
   ff = open(p_file,'r')
   count = 0
   for ln in ff:
       count += 1
       ln1 = ln.split('\n')
       ln2 = ln1[0].split('\t')
       if count % 100 == 0:
          if (p_st_using.upper())[0] == 'Y':
             print(' --- Processing row:', count,'/',LineNo)
       if count<3:
          if (p_st_using.upper())[0] == 'Y':
             print(ln2[0:5],ln2[-5:],len(ln2))
         
         #print(ln)
       if count == 1:
         L_colNo = len(ln2)
         L_title = ln2
         for idx in range(1,len(ln2)):
             L_rets[idx] = 0.0
             L_rets2[idx] = 0
       else:
         for idx in range(1,len(ln2)):
             va = float(ln2[idx])
             va2 = va
             if p_st>0:
               va2 = 2**va
             if va>L_max:
               L_max = va
             if va<L_min:
               L_min = va
             L_rets[idx] += va2
             if va>0:
               L_rets2[idx] += 1
   ff.close()
   L_rowNo = count
   if (p_st_using.upper())[0] == 'Y':
      print('Row#:',L_rowNo, ' Col#:',L_colNo)
      print('Min:',L_min, ' Max:',L_max)
   for idx in range(1,10):
      if (p_st_using.upper())[0] == 'Y':
         print(idx, L_rets[idx])
   return L_rets, L_rets2, L_title


def check_counts_in_sample_type(p_exp_count, p_cid, p_cid_sp_tp):
   L_rets = {}; L_ctype_all = {}; L_cell_tsize = {}
      
   for idx in p_exp_count:
      cid = p_cid[idx]
      t_exp_v = p_exp_count[idx]
      L_cell_tsize[cid] = t_exp_v
      
      cell_sample = p_cid_sp_tp[cid][0]
      cell_type = p_cid_sp_tp[cid][1]
      L_ctype_all[cell_type] = 1
      if not cell_sample in L_rets:
         L_rets[cell_sample] = {}
      if not cell_type in L_rets[cell_sample]:
         L_rets[cell_sample][cell_type] = []
      L_rets[cell_sample][cell_type].append(t_exp_v)

   return L_rets, L_ctype_all, L_cell_tsize


def get_sample_cell_type_exp_mean(p_sp_tp_count, p_fn_out, p_fn_out_2,p_ctypeAll, p_cell_count_LB):
   fout = open(p_fn_out,'w')
   L_tp_list = []
   for tp in p_ctypeAll:
      L_tp_list.append(tp)
   L_tp_list.sort()
   L_sp_list = []
   for sp in p_sp_tp_count:
      L_sp_list.append(sp)
   L_sp_list.sort()
   output = 'sample/cell_type\t' +'\t'.join(L_tp_list)+'\n'
   fout.write(output)
   for sp in L_sp_list:
      output = sp
      for tp in L_tp_list:
         if tp in p_sp_tp_count[sp]:
            t_values = p_sp_tp_count[sp][tp]
            if len(t_values)>=p_cell_count_LB:
               t_mean = numpy.mean(t_values)
               output += '\t'+str(t_mean)
            else:
               output += '\tnan'
         else:
            output += '\tnan'
      output += '\n'
      fout.write(output)
   fout.close()

   fout = open(p_fn_out_2,'w')
   output = 'sample/cell_type\t' +'\t'.join(L_tp_list)+'\n'
   fout.write(output)
   for sp in L_sp_list:
      output = sp
      for tp in L_tp_list:
         if tp in p_sp_tp_count[sp]:
            t_values_No = len(p_sp_tp_count[sp][tp])
            output += '\t'+str(t_values_No)
         else:
            output += '\t0'
      output += '\n'
      fout.write(output)  
   fout.close()


def check_meta_and_scRNAseq_data(p_meta, p_exp):
   L_rets = 1
   
   ff = open(p_exp, 'r')
   L_title = []
   count = 0
   for ln in ff:
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      count += 1
      if count == 1:
         L_title = ln2
      else:
         if len(ln2)==len(L_title):
            break
         L_rets = -1
         output = '\n !!! In the expression file, the numbers of columns in the first and second rows are not equal ---'
         output += '\n   --- Number of columns in the fist row: ' + str(len(L_title))
         output += '\n   --- Number of columns in the second row: ' + str(len(ln2))
         output += '\n !!! Please check the input data files before run the program --- '
                  

         print(output)
         break
   ff.close()
   if L_rets<0:
      return L_rets
   
   L_cell_all = {}; L_cell_list = []; L_sample_all = {}
   ff = open(p_meta, 'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln1 = ln.split('\t')
         L_cell_all[ln1[0]] = 1
         L_cell_list.append(ln1[0])
         if len(ln1)>2:
             L_sample_all[ln1[2]] = 1            
   ff.close()
   
   if len(L_sample_all)<2:
       L_rets = -1
       output = ''
       if len(L_sample_all)==1:
         output += '\n !!! The scRNA-seq data only has one sample/subject ---'
         output += '\n !!! There is no need for the scRNA-seq data normalization ---'
       else:
         output += '\n !!! The meta data does not have the third column for the sample/subject information ---'
         output += '\n !!! There is not enough information for performing the scRNA-seq data normalization --- '

       output += '\n !!! Please check the input data files before run the program --- '
               
       print(output)
       return L_rets
   
   if len(L_cell_list)+1 == len(L_title):
      t_matched = 0
      for idx in range(1, len(L_title)):
         if L_title[idx] in L_cell_all:
            t_matched += 1
      if t_matched<len(L_cell_all):
         L_rets = -1
         output = '\n !!! Some cell IDs in the meta and expression data files are not matched ---'
         output += '\n   --- Number of cells in the expression data file: ' + str(len(L_title)-1)
         output += '\n   --- Number of cells in the meta data file: ' + str(len(L_cell_all))
         output += '\n   --- Number of cells mathced in both files: ' + str(t_matched)
         output += '\n !!! Please check the input data files before run the program --- '
                  

         print(output)         
   else:
         L_rets = -1
         output = '\n !!! The numbers of cells in the meta and expression data files are not equal ---'
         output += '\n   --- Number of cells in the expression data file: ' + str(len(L_title)-1)
         output += '\n   --- Number of cells in the meta data file: ' + str(len(L_cell_list))
         output += '\n !!! Please check the input data files before run the program --- '
                  
         print(output)
   return L_rets

def get_sample_cell_type_exp_mean_and_cell_count(p_meta, p_exp, p_out_mean, p_out_count, p_out_cell_size, p_cell_count_LB = 10):
   L_cell_types =  read_meta_info_sample_ctype(p_meta)
   L_cell_read_counts, L_cell_expressed_gene_counts, L_title = check_scRNA(p_exp)
   L_count_cell_sample_types, L_all_types, L_cell_transcriptome_size = check_counts_in_sample_type(L_cell_read_counts,L_title,L_cell_types)
   get_sample_cell_type_exp_mean(L_count_cell_sample_types, p_out_mean, p_out_count, L_all_types, p_cell_count_LB)
   
   fout = open(p_out_cell_size, 'w')
   for cid in L_cell_transcriptome_size:
      output = cid+'\t'+str(L_cell_transcriptome_size[cid])+'\n'
      fout.write(output)
   fout.close()

def read_sp_tp_mean(p_file):
   L_rets = {}; L_gid_idx = {}
   ff = open(p_file,'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln1 = ln.split('\n')
         ln2 = ln1[0].split('\t')
         gid = ln2[0]
         L_rets[gid] = []
         L_gid_idx[count-2] = gid
         for idx in  range(1,len(ln2)):
            item = ln2[idx]
            if item == 'nan':
               L_rets[gid].append(-10000)
            else:
               L_rets[gid].append(float(item))
   ff.close()
   # print('----------------------')
   # print('Index\t\tSample_name')
   # for idx in L_gid_idx:
   #    print(' ',idx,'\t\t',L_gid_idx[idx])

   return L_rets


def get_means_two_samples(p_mean_all, p_sp1, p_sp2):
   L_sp_list_all = []
   for sp in p_mean_all:
      L_sp_list_all.append(sp)
   L_sp_list_all.sort()
   

   L_sp1 = L_sp_list_all[p_sp1]
   L_sp2 = L_sp_list_all[p_sp2]

   if (p_st_using.upper())[0] == 'Y':   
      print(' --- Draw grap for samples:', L_sp1, 'and', L_sp2)

   L_data_1 = p_mean_all[L_sp1]
   L_data_2 = p_mean_all[L_sp2]

   Y_list = []
   X_list = []
   for idx in range(len(L_data_1)):
       v1 = L_data_1[idx]
       v2 = L_data_2[idx]
       if v1>=0 and v2>=0:
         X_list.append(v1)
         Y_list.append(v2)
   return X_list, Y_list

def get_LR(p_sp_tp_mean, p_chosen_s, p_fix_idx = 0):
   L_rets = {}
   L_fixed_sp = ''
   count = -1
   L_sp_list = []
   for sp in p_sp_tp_mean:
      L_sp_list.append(sp)
   L_sp_list.sort()
   for sp in L_sp_list:
      count += 1
      if count == p_fix_idx:
         L_fixed_sp = sp
         break
   if (p_st_using.upper())[0] == 'Y':
      print('---Fixed sp:', L_fixed_sp)
   
   L_data_1 = p_sp_tp_mean[L_fixed_sp]
   for sp in p_chosen_s:
      L_data_2 = p_sp_tp_mean[sp]
      

      Y_list = []
      X_list = []
      X_tp = []
      for idx in range(len(L_data_1)):
          v1 = L_data_1[idx]
          v2 = L_data_2[idx]
          if v1>0 and v2>0:
            X_list.append([v1,1])
            Y_list.append(v2)
            X_tp.append(v1)
            
      # print('\n---+++')
      # print(L_fixed_sp, (X_tp))
      # print(sp, (Y_list))      
      
      x_train = numpy.array(X_list, dtype=numpy.float32)
      y_train = numpy.array(Y_list, dtype=numpy.float32)
   
      x_train_t = x_train.transpose()
      x_train_mul = numpy.dot(x_train_t,x_train)   
      x_train_inv = inv(x_train_mul)   
      x_result_ma = numpy.dot(x_train_inv, x_train_t)   
      L_para = numpy.dot(x_result_ma, y_train)
      
      L_rets[sp] = [L_para[0],L_para[1]]
      # print('a=',L_para[0], 'b=', L_para[1])
   return L_rets

def scRNA_seq_normalization(p_exp, p_exp_out, p_meta, p_sp_ctp_tra_size_mean, p_cell_size, p_chosen_samples, p_fixed_sample = 0):
   L_sp_tp_mean = read_sp_tp_mean(p_sp_ctp_tra_size_mean)
   L_Para_LR = get_LR(L_sp_tp_mean, p_chosen_samples, p_fixed_sample)
    
   L_v_max = -100; L_a_min = 100
   count = 0
   for sp in L_Para_LR:
     # if (p_st_using.upper())[0] == 'Y':
     #    print(sp,'\ta =',L_Para_LR[sp][0],'\tb =',L_Para_LR[sp][1],'\tratio =',L_Para_LR[sp][1]/L_Para_LR[sp][0])
     if L_Para_LR[sp][0] < L_a_min:
         L_a_min = L_Para_LR[sp][0]
     count += 1
     if count == 1:
         L_v_max = L_Para_LR[sp][1]/L_Para_LR[sp][0]
     else:
         t_ratio = L_Para_LR[sp][1]/L_Para_LR[sp][0]
         if L_v_max<t_ratio:
           L_v_max = t_ratio
     #print('--max ratio: ',L_v_max)
   if L_a_min<=0:
     if (p_st_using.upper())[0] == 'Y':

         print('!!! As at least one of constant "a" is negative, transcriptome sizes of cell types in at least two samples have problems.\n')
         print('    It is likely the cell type classification is not purified enough. The normalization will not perform.')
     return 0
   
   ff = open(p_exp,'r')
   LineNo = 0
   for ln in ff:
     LineNo += 1
   ff.close()
   
   LineNo2 = LineNo
   #print('---- Total row#:', LineNo)  
   
   L_Para_LR_2 = {}
   for sp in L_Para_LR:
     L_Para_LR_2[sp] = [L_Para_LR[sp][0],L_Para_LR[sp][1],(L_v_max-L_Para_LR[sp][1]/L_Para_LR[sp][0])/LineNo2]
     #print('  ', sp, [L_Para_LR[sp][0],L_Para_LR[sp][1],(L_v_max-L_Para_LR[sp][1]/L_Para_LR[sp][0])/LineNo])
    
   L_cell_types =  read_meta_info_sample_ctype(p_meta)
   count = 0
   for cid in L_cell_types:
     count += 1
     # if count<20:
     #    if (p_st_using.upper())[0] == 'Y':
     #       print(cid, L_cell_types[cid], L_Para_LR_2[ L_cell_types[cid][0]])
       
   fout = open(p_exp_out, 'w')
   ff = open(p_exp,'r')
   count = 0
   L_title = []
   for ln in ff:
     count += 1
     ln1 = ln.split('\n')
     ln2 = ln1[0].split('\t')
     if count%100 == 1:
        if (p_st_using.upper())[0] == 'Y':
           print('--- Processing row:',count,'/',LineNo, 'for normlaization......')
     if count == 1:
       L_title = ln2
       fout.write(ln)
     else:
       t_list = [ln2[0]]
       for idx in range(1, len(ln2)):
           t_cid = L_title[idx]
           
           t_cell_size = p_cell_size[t_cid]
           
           t_spid = L_cell_types[t_cid][0]
           t_v_a = L_Para_LR_2[t_spid][0]
           t_shift = L_Para_LR_2[t_spid][2]
           t_v_org = float(ln2[idx])
           
           #Add shift to all genes
           t_v_new = t_v_org/t_v_a + t_shift
           t_list.append(str(t_v_new))
           
           #Add shift to genes with positive expression
           # if t_cell_size>0:
           #    t_v_new = t_v_org/t_v_a*(1 + t_shift*t_v_a/t_cell_size)              
           #    t_list.append(str(t_v_new))
           # else:
           #    t_list.append(str(t_v_org))
       output = '\t'.join(t_list)+'\n'
       fout.write(output)
   ff.close()
   fout.close()
         

def get_LR_All(p_sp_tp_mean, p_fix_idx = 0):
   L_rets = {}
   L_fixed_sp = ''
   count = -1
   L_sp_list = []
   for sp in p_sp_tp_mean:
      L_sp_list.append(sp)
   L_sp_list.sort()
   for sp in L_sp_list:
      count += 1
      if count == p_fix_idx:
         L_fixed_sp = sp
         break
   if (p_st_using.upper())[0] == 'Y':
      print('---Fixed sp:', L_fixed_sp)
   
   L_data_1 = p_sp_tp_mean[L_fixed_sp]
   for sp in p_sp_tp_mean:
      L_data_2 = p_sp_tp_mean[sp]
      

      Y_list = []
      X_list = []
      X_tp = []
      for idx in range(len(L_data_1)):
          v1 = L_data_1[idx]
          v2 = L_data_2[idx]
          if v1>0 and v2>0:
            X_list.append([v1,1])
            Y_list.append(v2)
            X_tp.append(v1)
            
      # print('\n---+++')
      # print(L_fixed_sp, (X_tp))
      # print(sp, (Y_list))      
      
      x_train = numpy.array(X_list, dtype=numpy.float32)
      y_train = numpy.array(Y_list, dtype=numpy.float32)
   
      x_train_t = x_train.transpose()
      x_train_mul = numpy.dot(x_train_t,x_train)   
      x_train_inv = inv(x_train_mul)   
      x_result_ma = numpy.dot(x_train_inv, x_train_t)   
      L_para = numpy.dot(x_result_ma, y_train)
      
      L_rets[sp] = [L_para[0],L_para[1]]
      # print('a=',L_para[0], 'b=', L_para[1])
   return L_rets

def scRNA_seq_normalization_All(p_exp, p_exp_out, p_meta, p_sp_ctp_tra_size_mean, p_cell_size, p_fixed_sample = 0):
   L_sp_tp_mean = read_sp_tp_mean(p_sp_ctp_tra_size_mean)
   L_Para_LR = get_LR_All(L_sp_tp_mean, p_fixed_sample)
    
   L_v_max = -100; L_a_min = 100
   count = 0
   for sp in L_Para_LR:
     # if (p_st_using.upper())[0] == 'Y':
     #    print(sp,'\ta =',L_Para_LR[sp][0],'\tb =',L_Para_LR[sp][1],'\tratio =',L_Para_LR[sp][1]/L_Para_LR[sp][0])
     if L_Para_LR[sp][0] < L_a_min:
         L_a_min = L_Para_LR[sp][0]
     count += 1
     if count == 1:
         L_v_max = L_Para_LR[sp][1]/L_Para_LR[sp][0]
     else:
         t_ratio = L_Para_LR[sp][1]/L_Para_LR[sp][0]
         if L_v_max<t_ratio:
           L_v_max = t_ratio
     #print('--max ratio: ',L_v_max)
   if L_a_min<=0:
     if (p_st_using.upper())[0] == 'Y':

         print('!!! As at least one of constant "a" is negative, transcriptome sizes of cell types in at least two samples have problems.\n')
         print('    It is likely the cell type classification is not purified enough. The normalization will not perform.')
     return 0
   
   ff = open(p_exp,'r')
   LineNo = 0
   for ln in ff:
     LineNo += 1
   ff.close()
   LineNo2 = LineNo
   #print('---- Total row#:', LineNo)  
   
   L_Para_LR_2 = {}
   for sp in L_Para_LR:
     L_Para_LR_2[sp] = [L_Para_LR[sp][0],L_Para_LR[sp][1],(L_v_max-L_Para_LR[sp][1]/L_Para_LR[sp][0])/LineNo2]
     #print('  ', sp, [L_Para_LR[sp][0],L_Para_LR[sp][1],(L_v_max-L_Para_LR[sp][1]/L_Para_LR[sp][0])/LineNo])
    
   L_cell_types =  read_meta_info_sample_ctype(p_meta)
   count = 0
   for cid in L_cell_types:
     count += 1
     # if count<20:
     #    if (p_st_using.upper())[0] == 'Y':
     #       print(cid, L_cell_types[cid], L_Para_LR_2[ L_cell_types[cid][0]])
       
   fout = open(p_exp_out, 'w')
   ff = open(p_exp,'r')
   count = 0
   L_title = []
   for ln in ff:
     count += 1
     ln1 = ln.split('\n')
     ln2 = ln1[0].split('\t')
     if count%100 == 1:
        if (p_st_using.upper())[0] == 'Y':
           print('--- Processing row:',count,'/',LineNo, 'for normlaization......')
     if count == 1:
       L_title = ln2
       fout.write(ln)
     else:
       t_list = [ln2[0]]
       for idx in range(1, len(ln2)):
           t_cid = L_title[idx]
           
           t_cell_size = p_cell_size[t_cid]
           
           t_spid = L_cell_types[t_cid][0]
           t_v_a = L_Para_LR_2[t_spid][0]
           t_shift = L_Para_LR_2[t_spid][2]
           t_v_org = float(ln2[idx])

           #Add shift to all genes
           t_v_new = t_v_org/t_v_a + t_shift
           t_list.append(str(t_v_new))
           
           #Add shift to genes with positive expression
           # if t_cell_size>0:
           #    t_v_new = t_v_org/t_v_a*(1 + t_shift*t_v_a/t_cell_size)
           #    t_list.append(str(t_v_new))
           # else:
           #    t_list.append(str(t_v_org))
           
       output = '\t'.join(t_list)+'\n'
       fout.write(output)
   ff.close()
   fout.close()
      
def get_means_two_samples_2(p_mean_all, p_sp1, p_sp2):

   L_data_1 = p_mean_all[p_sp1]
   L_data_2 = p_mean_all[p_sp2]

   Y_list = []
   X_list = []
   for idx in range(len(L_data_1)):
       v1 = L_data_1[idx]
       v2 = L_data_2[idx]
       if v1>0 and v2>0:
         X_list.append(v1)
         Y_list.append(v2)
   return X_list, Y_list

def get_Pearson_correlation_samples(p_sp_tp_mean, p_cor_lb=0.95):
   L_Pearson = {}
   L_sp_list = []
   L_Pearson_spNo = {}
   L_cor_lb = 0.95
   for sp in p_sp_tp_mean:
       L_sp_list.append(sp)
   L_sp_list.sort()
   for tsp1 in L_sp_list:
       L_Pearson[tsp1] = []
       L_Pearson_spNo[tsp1] = []
       for tsp2 in L_sp_list:
          L_v1, L_v2 =  get_means_two_samples_2(p_sp_tp_mean, tsp1, tsp2)
          if len(L_v1)>3:
            t_pc = numpy.corrcoef(L_v1, L_v2)[0][1]
            L_Pearson[tsp1].append(t_pc)
            if t_pc>=p_cor_lb:
               L_Pearson_spNo[tsp1].append(tsp2)
          else:
            L_Pearson[tsp1].append(0)
   # print('----------------------')
   # print('Index\t\tSample_name\tStrong_correlation_Sample#')            
   # for idx in range(len(L_sp_list)):
   #    tsp = L_sp_list[idx]
   #    print(idx,'\t',tsp,'\t',len(L_Pearson_spNo[tsp]))
   #    print('   ---', L_Pearson_spNo[tsp])
   return L_Pearson, L_Pearson_spNo


def draw_heatmap(p_Pearson):
   df_corr = pd.DataFrame(p_Pearson)
   #df.columns = L_sp_list
   L_col_name = list(df_corr.columns)
   df_corr.index = L_col_name
   
   
   sns.light_palette("blue", as_cmap=True)
   #sns.color_palette("light:b", as_cmap=True)
   ax = sns.heatmap(df_corr, annot=False)
   
   return ax, df_corr


def draw_heatmap_Pearson_all(p_fn_ctyp_mean, p_fn_heatmap, p_fn_matrix):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, 0.95)
   #print(' --- Total sample#:', len(L_sp_tp_mean))

   L_figure, L_matrix = draw_heatmap(L_Pearson)
   L_heatmap = L_figure.get_figure()
   L_heatmap.savefig(p_fn_heatmap, dpi=400)
   #plt.savefig(p_fn_heatmap, dpi=400)
   
   L_matrix.to_csv(p_fn_matrix)
   
   
def draw_point(p_sp_tp_mean, p_idx1, p_idx2):
 
   L_X, L_Y = get_means_two_samples(p_sp_tp_mean, p_idx1, p_idx2)

   #draw the expression means of all cell types in two chosen samples 
   
   print('X:', L_X)
   print('Y:', L_Y)
   plt.plot(L_X, L_Y, 'o')
   plt.show()
   return L_X, L_Y

def draw_cell_type_size_mean_two_samples(p_fn_ctyp_mean):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   
   
   #Choose any two samples to draw point graph
   t_sp_1_idx = int(input("Input the index of sample1 (0-"+str(len(L_sp_tp_mean)-1)+"): "))
   t_sp_2_idx = int(input("Input the index of sample2 (0-"+str(len(L_sp_tp_mean)-1)+"): "))
   tX, tY = draw_point(L_sp_tp_mean, t_sp_1_idx, t_sp_2_idx)

def get_sample_baseline(p_fn_ctyp_mean, L_Pearson_LB):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, L_Pearson_LB)
   L_sp_list = []
   for sp in L_samples:
      L_sp_list.append(sp)
   L_sp_list.sort()
   
   L_rets = 0
   L_lenTp = len(L_samples[L_sp_list[0]])
   for idx in range(len(L_sp_list)):
      tsp = L_sp_list[idx]
      if L_lenTp<len(L_samples[tsp]):
         L_rets = idx
         L_lenTp = len(L_samples[tsp])
   return L_rets


def get_W0_W1_W(X, Y):
   X_sum = 0.0; Y_sum = 0.0
   XY = 0.0; X_squre = 0.0
   
   t_N = len(X)
   for idx in range(len(X)):
      t_x = X[idx]; t_y = Y[idx]
      X_sum += t_x
      Y_sum += t_y
      XY += t_x*t_y
      X_squre += t_x*t_x

   
   # t_W1 = (XY-X_sum*Y_sum/t_N)/(X_squre-X_sum*X_sum/t_N)
   # t_W0 = Y_sum/t_N-t_W1*X_sum/t_N
   
   t_W1 = (t_N*XY-X_sum*Y_sum)/(t_N*X_squre-X_sum*X_sum)
   t_W0 = (Y_sum-t_W1*X_sum)/t_N
   t_W = XY/X_squre


   x_min = numpy.min(X)
   x_max = numpy.max(X)

   
   L_P0 = [x_min, x_max]
   L_P1 = [t_W1*x_min+t_W0, t_W1*x_max+t_W0]
   L_P2 = [t_W*x_min, t_W*x_max]

   #print('--W1, W0, W, Wr', t_W1, t_W0, t_W, t_Wr)
   return L_P0, L_P1, L_P2

def draw_cell_type_size_mean_point_plot(p_fn_ctyp_mean, p_fn_point, L_base, L_Pearson_LB, p_gn_row=3, p_fl_sft = 1, p_fl_noSft = 0):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, L_Pearson_LB)
   L_sp_list = []
   for sp in L_samples:
      L_sp_list.append(sp)
   L_sp_list.sort()
   L_sp_base = L_sp_list[L_base]
   L_chosen_sp = L_samples[L_sp_base]
   #print(L_sp_base, L_chosen_sp)
   
   if len(L_chosen_sp)<2:
      print('\n---- Only one sample is chosen, no point plot figure!!!! \n')
      return 1
   
   count = 0
   colNo = p_gn_row
   rowNo = math.ceil((len(L_chosen_sp)-1)/colNo)
   
   if len(L_chosen_sp)>2:
   
      fig, axs = plt.subplots(rowNo, colNo, figsize=(6*colNo,4.5*rowNo))
      for t_sp2 in L_chosen_sp:
         if not t_sp2 == L_sp_base:
            L_v1, L_v2 =  get_means_two_samples_2(L_sp_tp_mean, L_sp_base, t_sp2)
            t_P0, t_P1, t_P2 = get_W0_W1_W(L_v1, L_v2)
            
            t_c_idx = count % colNo
            t_r_idx = math.floor(count/colNo)
            # print('\n-----',L_sp_base,t_sp2, t_r_idx, t_c_idx)
            # print(L_v1)
            # print(L_v2)
            #axs[t_r_idx,t_c_idx].plot(L_v1, L_v2, 'o')
            if colNo == 1:
               axs[t_r_idx].plot(L_v1, L_v2, 'o')
               if not p_fl_sft == 0:
                  axs[t_r_idx].plot(t_P0, t_P1, '--', color = 'green', label='With_shift')
               if not p_fl_noSft == 0:
                  axs[t_r_idx].plot(t_P0, t_P2, '--', color = 'red', label='No_shift')
                  
               axs[t_r_idx].set_xlabel(L_sp_base)
               axs[t_r_idx].set_ylabel(t_sp2)
               if p_fl_sft>0 or p_fl_noSft>0:
                  axs[t_r_idx].legend(title = 'Fit line')
            else:
               if rowNo == 1:
                  axs[t_c_idx].plot(L_v1, L_v2, 'o')
                  if not p_fl_sft == 0:
                     axs[t_c_idx].plot(t_P0, t_P1, '--', color = 'green', label='With_shift') 
                  if not p_fl_noSft == 0:
                     axs[t_c_idx].plot(t_P0, t_P2, '--', color = 'red', label='No_shift')
                     
                  axs[t_c_idx].set_xlabel(L_sp_base)
                  axs[t_c_idx].set_ylabel(t_sp2)
                  if p_fl_sft>0 or p_fl_noSft>0:
                     axs[t_c_idx].legend(title = 'Fit line')
               else:
                  axs[t_r_idx,t_c_idx].plot(L_v1, L_v2, 'o')
                  if not p_fl_sft == 0:
                     axs[t_r_idx,t_c_idx].plot(t_P0, t_P1, '--', color = 'green', label='With_shift') 
                  if not p_fl_noSft == 0:
                     axs[t_r_idx,t_c_idx].plot(t_P0, t_P2, '--', color = 'red', label='No_shift')
                     
                  axs[t_r_idx,t_c_idx].set_xlabel(L_sp_base)
                  axs[t_r_idx,t_c_idx].set_ylabel(t_sp2)
                  if p_fl_sft>0 or p_fl_noSft>0:
                     axs[t_r_idx,t_c_idx].legend(title = 'Fit line')
            count += 1
   
            
      fig.savefig(p_fn_point, dpi=400)
   else:
      if len(L_chosen_sp) == 2:
            for t_sp2 in L_chosen_sp:
               if not t_sp2 == L_sp_base:
                  L_v1, L_v2 =  get_means_two_samples_2(L_sp_tp_mean, L_sp_base, t_sp2)

                  t_P0, t_P1, t_P2 = get_W0_W1_W(L_v1, L_v2)                  
                            
                  figure = Figure(figsize=(8,7.5), dpi=100)
                  a = figure.add_subplot(111)
                  a.plot(L_v1, L_v2, 'o') 
                  if not p_fl_sft == 0:
                     a.plot(t_P0, t_P1, '--', color = 'green', label='With_shift')
                  if not p_fl_noSft == 0:
                     a.plot(t_P0, t_P2, '--', color = 'red', label='No_shift')
                  a.legend(title = 'Fit line')
                  
                  figure.savefig(p_fn_point, dpi=400)
                  
                  plt.plot(L_v1, L_v2, 'o') 
                  
                  if not p_fl_sft == 0:
                     plt.plot(t_P0, t_P1, '--', color = 'blue', label='With_shift') 
                  if not p_fl_noSft == 0:
                     plt.plot(t_P0, t_P2, '--', color = 'red', label='No_shift') 
                  if p_fl_sft>0 or p_fl_noSft>0:
                     plt.legend(title = 'Fit line')
                  plt.show()
                  
      
def get_sp_ctp_cellNo(p_sp_ctp, p_samples):
   L_rets = {}
   for sp in p_samples:
      for ctp in p_sp_ctp[sp]:
         if not ctp in L_rets:
            L_rets[ctp] = 0
         L_rets[ctp] += p_sp_ctp[sp][ctp]
   L_ctp_list = []
   for ctp in L_rets:
      L_ctp_list.append(ctp)
   L_ctp_list.sort()
   output1 = ''
   output2 = ''
   for ctp in L_ctp_list:
      output1 += '\t'+ctp
      output2 += '\t'+str(L_rets[ctp])
   return output1, output2
   
   
   
def get_cellType_cellNumber(p_meta, p_samples_Pearson):
   L_sample_cellType_map = {}
   ff = open(p_meta, 'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln1 = ln.split('\n')
         ln2 = ln1[0].split('\t')
         tctp = ln2[1]
         tspid = ln2[2]
         if not tspid in L_sample_cellType_map:
            L_sample_cellType_map[tspid] = {}
         if not tctp in L_sample_cellType_map[tspid]:
            L_sample_cellType_map[tspid][tctp] = 0
         L_sample_cellType_map[tspid][tctp] += 1
   ff.close()
   L_sp_list = []
   for sp in p_samples_Pearson:
      L_sp_list.append(sp)
   L_sp_list.sort()
   
   
   for idx in range(len(L_sp_list)):
      tsp = L_sp_list[idx]
      print('\n',idx,'\t',tsp,'\t',len(p_samples_Pearson[tsp]),p_samples_Pearson[tsp])
      t_output1, t_output2 = get_sp_ctp_cellNo(L_sample_cellType_map, p_samples_Pearson[tsp])
      print(t_output1)
      print(t_output2 )
      #print(p_samples_Pearson[tsp])


def get_sample_cell_type_information_top_Pearson(p_fn_meta, p_fn_ctyp_mean, p_Pearson_LB = 0.7):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, p_Pearson_LB)
   # print(' --- Total sample#:', len(L_sp_tp_mean))

   
   get_cellType_cellNumber(p_fn_meta, L_samples)   

def get_cellType_cellNumber_2(p_count, p_samples_Pearson):
   L_sample_cellType_map = {}
   
   L_rets = {}
   L_title = []
   ff = open(p_count, 'r')
   count = 0
   for ln in ff:
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      count += 1
      if count == 1:
         L_title = ln2
      else:
         tsp = ln2[0]
         L_sample_cellType_map[tsp] = {}
         for idx in range(1, len(ln2)):
            tpid = L_title[idx]
            t_cno = int(ln2[idx])
            L_sample_cellType_map[tsp][tpid] = t_cno
   ff.close()  
   
   L_sp_list = []
   for sp in p_samples_Pearson:
      L_sp_list.append(sp)
   L_sp_list.sort()
   
   output_All = ''
   for idx in range(len(L_sp_list)):
      tsp = L_sp_list[idx]
      #print('\n',idx,'\t',tsp,'\t',len(p_samples_Pearson[tsp]),p_samples_Pearson[tsp])
      output = '\n'+str(idx) +'\t' + tsp +'\t' + str(len(p_samples_Pearson[tsp])) + '\t'+ str(p_samples_Pearson[tsp]) +'\n'
      output_All += output
      
      
      t_output1, t_output2 = get_sp_ctp_cellNo(L_sample_cellType_map, p_samples_Pearson[tsp])
      #print(t_output1)
      output = t_output1+'\n'
      output_All += output
      #print(t_output2)
      output = t_output2+'\n'
      output_All += output
      #print(p_samples_Pearson[tsp])
   return output_All

def get_sample_cell_type_information_top_Pearson_2(p_fn_count, p_fn_ctyp_mean, p_fn_extra_info, p_Pearson_LB = 0.7):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, p_Pearson_LB)
   # print(' --- Total sample#:', len(L_sp_tp_mean))

   
   L_output_All = get_cellType_cellNumber_2(p_fn_count, L_samples)
   
   fout = open(p_fn_extra_info, 'w')
   fout.write(L_output_All)
   fout.close()  
   print(L_output_All)

   
def get_chosen_samples_heatmap(p_base, p_samples_Pearson, p_threshold, p_sp_ctp_mean):
   L_sp_list = []
   L_rets = []; L_sp_ctp_mean2 = {}
   for sp in p_samples_Pearson:
      L_sp_list.append(sp)
   L_sp_list.sort()
   L_sp_base = L_sp_list[p_base]
   for idx in range(len(p_samples_Pearson[L_sp_base])):
      tva = p_samples_Pearson[L_sp_base][idx]
      if tva>=p_threshold:
         L_rets.append(L_sp_list[idx])
         
   for sp in L_rets:
      L_sp_ctp_mean2[sp] = p_sp_ctp_mean[sp]
   print('---Chose samples:', L_rets, len(L_rets))
   #print('***', p_sp_ctp_mean)
      
   L_Pearson2, L_samples2 = get_Pearson_correlation_samples(L_sp_ctp_mean2, p_threshold)
   draw_heatmap(L_Pearson2)
   
   return L_rets

def Make_small_meta_data_sample(p_fn_meta_small,p_fn_meta_small_malig, p_sp_chosen):
   fout = open(p_fn_meta_small_malig, 'w')
   output = 'Cell_ID\tCell_type\tSample_ID\n'
   fout.write(output)
   ff = open(p_fn_meta_small,'r')
   count = 0
   for ln in ff:
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      ctp = ln2[2]
      if ctp in p_sp_chosen:
         fout.write(ln)
         count += 1

   ff.close()
   fout.close()
   print('--Total chosen cell#:', count)
      
def get_sub_matrix_sp_col_2(p_fn_exp, p_fn_meta, p_fn_out):
   #The order of columnn in the submatrix is the order of row in mata data
   L_sp_chosen = {}; L_sp_order = []
   L_cellID_idx_map = {}
   ff = open(p_fn_meta,'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln1 = ln.split('\t')
         L_sp_chosen[ln1[0]] = 1
         L_sp_order.append(ln1[0])
   ff.close()
   #print('--- Total cell# info read:', count-1)
   fout = open(p_fn_out,'w')
   ff = open(p_fn_exp,'r')
   count = 0; count2 = 0
   #L_idx_in = []
   for ln in ff:
      count += 1
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      if count%100 == 1:
         print('---Procesing row', count, ' for new matrix ......')
      if count == 1:
         output = 'Gene_sample'
         for idx in range(1,len(ln2)):
            tcid = ln2[idx]
            L_cellID_idx_map[tcid] = idx
         for tcid in L_sp_order:
            output += '\t'+tcid
            #L_idx_in.append(idx)
         output += '\n'
         fout.write(output)
      else:
         output = ln2[0]
         for tcid in L_sp_order:
            tidx = L_cellID_idx_map[tcid]
            output += '\t'+ln2[tidx]
         output += '\n'
         fout.write(output)
            
         

            
   ff.close()
   fout.close()
   #print('---- total cell#:', len(L_idx_in))   

def get_meta_and_exp_chosen_sample(p_exp, p_meta, p_exp_out, p_meta_out, p_chosen_sp):
   L_sp_dic = {}
   for sp in p_chosen_sp:
      L_sp_dic[sp] = 1
   Make_small_meta_data_sample(p_meta, p_meta_out, L_sp_dic)
   
   print('--- Making the new scRNA-seq data matrix for chosen samples ---')
   get_sub_matrix_sp_col_2(p_exp, p_meta_out, p_exp_out)  
   
      
def get_and_normalize_cell_subset(p_fn_exp, p_fn_meta, p_fn_ctyp_mean, p_fn_cell_size, p_fn_exp_2, p_fn_meta_2, p_fn_exp_3, p_L_sample_based, p_L_Pearson_LB):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, p_L_Pearson_LB)


   L_sample1 = get_chosen_samples_heatmap(p_L_sample_based, L_Pearson, p_L_Pearson_LB, L_sp_tp_mean)

   print('--- Get cell subset......')
   get_meta_and_exp_chosen_sample(p_fn_exp, p_fn_meta, p_fn_exp_2, p_fn_meta_2, L_sample1)
   
   print('--- Normalization.....')
   L_cell_size = read_cell_transcriptome_size(p_fn_cell_size)
   scRNA_seq_normalization(p_fn_exp_2, p_fn_exp_3, p_fn_meta_2, p_fn_ctyp_mean, L_cell_size, L_sample1,  p_L_sample_based)

def read_cell_transcriptome_size(p_file):
   L_rets = {}
   ff = open(p_file, 'r')
   for ln in ff:
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      cid = ln2[0]
      tsize = float(ln2[1])
      L_rets[cid] = tsize
   ff.close()
   return L_rets

def get_cell_subset_scRNA_seq_data_normalization(p_fn_exp, p_fn_meta, p_fn_ctyp_mean, p_fn_cell_size, p_fn_exp_2, p_fn_meta_2, p_fn_exp_3, p_L_sample_based, p_L_Pearson_LB):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, p_L_Pearson_LB)


   L_sample1 = get_chosen_samples_heatmap(p_L_sample_based, L_Pearson, p_L_Pearson_LB, L_sp_tp_mean)
   
   print(len(L_sp_tp_mean), len(L_samples), len(L_sample1))
   
   L_cell_size = read_cell_transcriptome_size(p_fn_cell_size)
   
   if len(L_sample1) == len(L_sp_tp_mean):
      shutil.copyfile(p_fn_meta, p_fn_meta_2)
      if len(L_sample1)>1:
         print('--- Do normalization.....')
         scRNA_seq_normalization_All(p_fn_exp, p_fn_exp_3, p_fn_meta, p_fn_ctyp_mean, L_cell_size,  p_L_sample_based)
      else:
         output = '\n!!!! Only one sample is chosen. No need for doing normalization ------'
         shutil.copyfile(p_fn_exp, p_fn_exp_3)
         print(output)
   else:
      
      print('--- Get cell subset......')
      get_meta_and_exp_chosen_sample(p_fn_exp, p_fn_meta, p_fn_exp_2, p_fn_meta_2, L_sample1)
      
      if len(L_sample1)>1:
         print('--- Do normalization.....')
         scRNA_seq_normalization(p_fn_exp_2, p_fn_exp_3, p_fn_meta_2, p_fn_ctyp_mean, L_cell_size, L_sample1,  p_L_sample_based)
      else:
         output = '\n!!!! Only one sample is chosen. No need for doing normalization ------'
         shutil.copyfile(p_fn_exp_2, p_fn_exp_3)
         print(output)



#---Normalization without shift

def read_meta_Sample_CellType_cid(p_fn_meta, p_cid=0, p_ctp_clust=1, p_sp=2):
   L_rets = {}; L_ctp_clu_all = {}; L_cid = {}
   ff = open(p_fn_meta, 'r')
   count = 0
   for ln in ff:
      count += 1
      if count>1:
         ln1 = ln.split('\n')
         ln2 = ln1[0].split('\t')
         cid = ln2[p_cid]
         ctp = ln2[p_ctp_clust]
         tsp = ln2[p_sp]
         if not tsp in L_rets:
            L_rets[tsp] = {}
         if not ctp in L_rets[tsp]:
            L_rets[tsp][ctp] = {}
         L_rets[tsp][ctp][cid] = 1
         L_ctp_clu_all[ctp] = 1
         
         L_cid[cid] = count
         
   ff.close()
   
   if testShow=='Y':
      print('\nSample cell number in each cell type ==========================')
      for tsp in L_rets:
         print(tsp, '-------------')
         for ctp in L_rets[tsp]:
            print(' ', ctp,'\t', len(L_rets[tsp][ctp]))
            
   print('====', len(L_cid), count)
            
   return L_rets, L_ctp_clu_all

def get_CellType_Sample_ftSize_CellNo(p_fn_tfsz, p_meta):
   L_rets = {}
   ff = open(p_fn_tfsz, 'r')
   count = 0
   L_title = []
   for ln in ff:
      count += 1
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      if count == 1:
         L_title = ln2
      else:
         #ctp = ln2[0]
         tsp = ln2[0]
         for idx in range(1, len(ln2)):
            #tsp = L_title[idx]
            ctp = L_title[idx]
            tva = ln2[idx]
            if not tva == 'nan':
               tva2 = float(tva)
               tcellNO = len(p_meta[tsp][ctp])
               if not ctp in L_rets:
                  L_rets[ctp] = {}
               L_rets[ctp][tsp] = [tcellNO, tva2]
   ff.close()
   if testShow=='Y':
      for ctp in L_rets:
         print(ctp, '-----')
         for sp in L_rets[ctp]:
            print(' ', sp, L_rets[ctp][sp])
   
   return L_rets

def get_normalization_ratio_baseLine_and_anotherSample_simple_mean(p_ctp_ftSize_CellNo, p_baseLine_sp, p_sample_2):
   #w is to minimize  = (y1-wx1)^2 + (y2-wx2)^2 + ... ; adjust ratio = 1/w; linear regression with w0=0
   L_rets = -1
   L_ratio_all = []; L_size_bl = []; L_size_anotherOne = []
   if testShow=='Y':
      print('Samples:', p_baseLine_sp, p_sample_2)
   L_deno = 0.0; L_neum = 0.0
   for ctp in p_ctp_ftSize_CellNo:
      if p_baseLine_sp in p_ctp_ftSize_CellNo[ctp] and p_sample_2 in p_ctp_ftSize_CellNo[ctp]:
         tsz_bl = p_ctp_ftSize_CellNo[ctp][p_baseLine_sp][1]
         tsz_other_sp = p_ctp_ftSize_CellNo[ctp][p_sample_2][1]
         # tRatio = tsz_bl/tsz_other_sp
         # L_ratio_all.append(tRatio)
         L_deno += tsz_bl*tsz_bl
         L_neum += tsz_other_sp*tsz_bl
         L_size_bl.append(tsz_bl)
         L_size_anotherOne.append(tsz_other_sp)
         if testShow=='Y':
            print(' ', ctp, '\t', tsz_bl, '\t', tsz_other_sp)
            
   #L_rets = L_deno/L_neum #Causeing "float division by zero" if the baseline sample and any other sample do not share any cell type. Will be changed.
   #--- Modifying begins.
   L_rets = 0.0
   if L_neum > 0:
       L_rets = L_deno/L_neum
   #--- Modifying ends. 03/22/2025.
   
   L_Pearson_coef = 0
   L_size_bl.append(0.0); L_size_anotherOne.append(0.0)
   if len(L_size_bl)>=3:
      L_Pearson_coef = t_pc = numpy.corrcoef(L_size_bl, L_size_anotherOne)[0][1]
   if testShow=='Y':
      print('Ratio mean:', L_rets, ' Pearson Coef:', L_Pearson_coef)
   return L_rets

def get_normalization_ratio_sample_mean(p_ctp_ftSize_CellNo, p_sp_ctp_cid, p_base_line_sample):
   L_rets = {}
   L_sp_list = []
   for sp in p_sp_ctp_cid:
      L_sp_list.append(sp)
   L_sp_list.sort()
   
   L_baseLine_sp = L_sp_list[p_base_line_sample]
   
   for idx in range(len(L_sp_list)):
      if not idx == p_base_line_sample:
         t_sample_2 = L_sp_list[idx]
         t_ratio = get_normalization_ratio_baseLine_and_anotherSample_simple_mean(p_ctp_ftSize_CellNo, L_baseLine_sp, t_sample_2)
         L_rets[t_sample_2] = t_ratio
   L_rets[L_baseLine_sp] = 1.0 

         
   return L_rets

def get_normalization_ratio_cell_type_sp(p_ctp_ftSize_CellNo, p_sp_ratio, p_adjust_status = 0):
   #1. Compute the transcriptome size of each cell type: weight sum of size in each sample. Weight of sp1 = (sp1_cell#)/(sp1_cell# + sp2_cell#+...)
   #2. Get adjust of each cell type and each sample.
   #3. If p_adjust_status ==0, not adjust transcriptome size of the same cell type in different samples to the same value.
   L_rets = {}
   if p_adjust_status == 0:
      for ctp in p_ctp_ftSize_CellNo:
         L_rets[ctp] = {}
         for sp in p_ctp_ftSize_CellNo[ctp]:
             L_rets[ctp][sp] = p_sp_ratio[sp]
   else:
      for ctp in p_ctp_ftSize_CellNo:
         L_rets[ctp] = {}
         #print(ctp, '===========')
         tsum_fts = 0.0
         t_cNo_all = 0.0
         for sp in p_ctp_ftSize_CellNo[ctp]:
             tratio = p_sp_ratio[sp]
             #print('  ', sp, p_ctp_ftSize_CellNo[ctp][sp], tratio, tratio*p_ctp_ftSize_CellNo[ctp][sp][1])
             t_cNo = p_ctp_ftSize_CellNo[ctp][sp][0]
             tsum_fts += p_ctp_ftSize_CellNo[ctp][sp][1]*tratio*t_cNo
             t_cNo_all += t_cNo
         tsum_fts_avg = tsum_fts/t_cNo_all
         #print('  ** new avg ratio:', tsum_fts_avg)
         for sp in p_ctp_ftSize_CellNo[ctp]:
             t_fts = p_ctp_ftSize_CellNo[ctp][sp][1]
             tratio_new = tsum_fts_avg/t_fts
             #print(' --Ratio_new, Ratio_old', sp, tratio_new, p_sp_ratio[sp])
             L_rets[ctp][sp] = tratio_new
         
   return L_rets

def get_cid_normalization_ratio(p_normalization_ratio_CellType_sp, p_normalization_ratio_sp, p_sp_ctp_cid):
   #Get the ajust ratio for each cell
   L_rets = {}
   for sp in p_sp_ctp_cid:
      for ctp in p_sp_ctp_cid[sp]:
         if ctp in p_normalization_ratio_CellType_sp:
            if sp in p_normalization_ratio_CellType_sp[ctp]:
               tRatio = p_normalization_ratio_CellType_sp[ctp][sp]
               for cid in p_sp_ctp_cid[sp][ctp]:
                  L_rets[cid] = tRatio
            else:
               tRatio = p_normalization_ratio_sp[sp]
               for cid in p_sp_ctp_cid[sp][ctp]:
                   L_rets[cid] = tRatio                
         else:
            tRatio = p_normalization_ratio_sp[sp]
            for cid in p_sp_ctp_cid[sp][ctp]:
                L_rets[cid] = tRatio   
               
   return L_rets

def get_cellType_sample_CellNo_and_tSizeMean(p_fn_ctp_ftsz, p_fn_meta, p_base_line_sample = 0):
   
   L_sp_ctp_cid, L_ctp_info = read_meta_Sample_CellType_cid(p_fn_meta)
   
   
   L_ctp_ftSize_CellNo = get_CellType_Sample_ftSize_CellNo(p_fn_ctp_ftsz, L_sp_ctp_cid)
   
   L_normalization_ratio_sp = get_normalization_ratio_sample_mean(L_ctp_ftSize_CellNo, L_sp_ctp_cid, p_base_line_sample)
   
   # print('======')
   # for sp in L_normalization_ratio_sp:
   #     print(sp, L_normalization_ratio_sp[sp])
      
   L_normalization_ratio_CellType_sp = get_normalization_ratio_cell_type_sp(L_ctp_ftSize_CellNo, L_normalization_ratio_sp)

   # print('======')
   # for sp in L_normalization_ratio_sp:
   #     print(' old',sp, L_normalization_ratio_sp[sp], len(L_ctp_ftSize_CellNo))
      
   L_cid_ratio = get_cid_normalization_ratio(L_normalization_ratio_CellType_sp, L_normalization_ratio_sp, L_sp_ctp_cid)
   
   #print(' *** Cell Number:', len(L_cid_ratio))
   
   return L_cid_ratio


def do_normalization(p_fn_scRNAseq_old, p_fn_scRNAseq_CBTS, p_cid_ratio):
   L_idx_cid_ratio = {}; L_idx_list = []
   fout = open(p_fn_scRNAseq_CBTS, 'w')
   ff = open(p_fn_scRNAseq_old, 'r')
   count = 0
   for ln in ff:
      count += 1
      ln1 = ln.split('\n')
      ln2 = ln1[0].split('\t')
      if count%100 == 1:
         print('--- Normalizing row', count)
      if count == 1:
         cell_chosen = [ln2[0]]
         for idx in range(1, len(ln2)):
            item = ln2[idx]
            if item in p_cid_ratio:
               L_idx_cid_ratio[idx] = p_cid_ratio[item]
               L_idx_list.append(idx)
               cell_chosen.append(item)
         output = '\t'.join(cell_chosen)
         fout.write(output)
         fout.write('\n')
      else:
         vas_all = [ln2[0]]
         for idx in L_idx_list:
            tva_old = float(ln2[idx])
            t_ratio = L_idx_cid_ratio[idx]
            tva_new = str(tva_old*t_ratio)
            vas_all.append(tva_new)
         output = '\t'.join(vas_all)
         fout.write(output)
         fout.write('\n')  

   ff.close()
   fout.close()  

def get_cell_subset_scRNA_seq_data_normalization_no_shift(p_fn_exp, p_fn_meta, p_fn_ctyp_mean, p_fn_cell_size, p_fn_exp_2, p_fn_meta_2, p_fn_exp_3, p_L_sample_based, p_L_Pearson_LB):
   L_sp_tp_mean = read_sp_tp_mean(p_fn_ctyp_mean)
   L_Pearson, L_samples = get_Pearson_correlation_samples(L_sp_tp_mean, p_L_Pearson_LB)


   L_sample1 = get_chosen_samples_heatmap(p_L_sample_based, L_Pearson, p_L_Pearson_LB, L_sp_tp_mean)
   
   print(len(L_sp_tp_mean), len(L_samples), len(L_sample1))
   
   L_ctp_ftsz_cNo = get_cellType_sample_CellNo_and_tSizeMean(p_fn_ctyp_mean, p_fn_meta, p_L_sample_based)
   
   # L_cell_size = read_cell_transcriptome_size(p_fn_cell_size)
   
   if len(L_sample1) == len(L_sp_tp_mean):
       shutil.copyfile(p_fn_meta, p_fn_meta_2)
       if len(L_sample1)>1:
         print('--- Do normalization.....')
         do_normalization(p_fn_exp, p_fn_exp_3,L_ctp_ftsz_cNo)
       else:
         output = '\n!!!! Only one sample is chosen. No need for doing normalization ------'
         shutil.copyfile(p_fn_exp, p_fn_exp_3)
         print(output)
   else:
      
       print('--- Get cell subset......')
       get_meta_and_exp_chosen_sample(p_fn_exp, p_fn_meta, p_fn_exp_2, p_fn_meta_2, L_sample1)
      
       if len(L_sample1)>1:
         print('--- Do normalization.....')
         do_normalization(p_fn_exp_2, p_fn_exp_3,L_ctp_ftsz_cNo)
       else:
         output = '\n!!!! Only one sample is chosen. No need for doing normalization ------'
         shutil.copyfile(p_fn_exp_2, p_fn_exp_3)
         print(output)





#--End of Normalization without shift