# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2024-01-05 20:28:47
# @Last Modified by:   dongqing
# @Last Modified time: 2024-04-17 03:34:46


import os
import warnings
# pd.set_option('mode.chained_assignment', None)
# warnings.filterwarnings('ignore', category=RuntimeWarning)
# warnings.filterwarnings('ignore', category=FutureWarning, append=True)
warnings.filterwarnings('ignore')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 

import spateo as st
import numpy as np

from skimage.io import imread, imsave
from Cellist.Utility import *
from Cellist.IO import *
from Cellist.Plot import *

import matplotlib
import matplotlib.pyplot as plt

plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica']

'''
cellist align --gem /mnt/Venus/home/dongxin/Projects/heart/202401_Stereo_seq/result_C02942B3/041.cellcut/C02942B3.adjusted.gem \
--tif /mnt/Venus/home/dongxin/Projects/heart/202401_Stereo_seq/result_C02942B3/03.register/ssDNA_C02942B3_regist.tif \
--nworkers 24 --outdir .
'''

def AlignmentParser(subparsers):
    parser = subparsers.add_parser("align", 
        help = "Refine alignment between image and spatial transcriptomics. ")
    group_input = parser.add_argument_group("Input arguments")
    group_input.add_argument("--gem", dest = "spot_expr_file", default = None,
        help = "Bin1 gem file generated by SAW pipeline located in '041.cellcut' or '04.tissuecut'. ")
    group_input.add_argument("--tif", dest = "regist_tif_file", default = None,
        help = "Registered tif file generated by SAW pipeline located in '04.tissuecut'. ")
    group_input.add_argument("--nworkers", dest = "num_workers", type = int, default = 8,
        help = "Maximum number of workers to use. ")

    group_output = parser.add_argument_group("Output arguments")
    group_output.add_argument("--outdir", dest = "out_dir", default = '.',
        help = "The output directory. ")
    group_output.add_argument("--outprefix", dest = "out_prefix", default = None,
        help = "The output prefix. ")
    return(parser)

def Alignment(gem_path, img_path, num_workers, out_dir, out_prefix):
    st.config.n_threads = num_workers
    if not out_prefix:
        file_name, file_extension = os.path.splitext(gem_path)
        base_name = os.path.basename(file_name)
        out_gem_prefix = base_name
        file_name, file_extension = os.path.splitext(img_path)
        base_name = os.path.basename(file_name)
        out_img_prefix = base_name
    else:
        out_gem_prefix = out_prefix
        out_img_prefix = out_prefix
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    gem_df = pd.read_csv(gem_path, sep = "\t", header = 0)
    x_range = gem_df['x'].max() -  gem_df['x'].min()
    y_range = gem_df['y'].max() -  gem_df['y'].min() 
    # temporary files
    gem_path_spliced = os.path.join(out_dir, "%s_spliced.gem" %out_gem_prefix)
    img = imread(img_path, as_gray = True)
    fig, ax = plt.subplots(1, 1, figsize=(20, 20), sharex=True, sharey=True)
    ax.imshow(img, cmap="gray")
    # ax.set_title("Original", fontsize = 20)
    # plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "%s_original.png" %out_img_prefix), bbox_inches='tight')
    # ratation
    if abs(img.shape[0] - x_range)*abs(img.shape[1] - y_range) > abs(img.shape[0] - y_range)*abs(img.shape[1] - x_range):
        img_path_rotated = os.path.join(out_dir, "%s_transposed.tif" %out_img_prefix)
        rotated = img.T
        rotated = rotated.astype(np.int32)
        imsave(img_path_rotated, rotated)
        fig, ax = plt.subplots(1, 1, figsize=(20, 20), sharex=True, sharey=True)
        ax.imshow(rotated, cmap="gray")
        # ax.set_title("Original", fontsize = 20)
        # plt.tight_layout()
        plt.savefig(os.path.join(out_dir, "%s_transposed.png" %out_img_prefix), bbox_inches='tight')
    else:
        img_path_rotated = img_path
        rotated = img.copy()
        rotated = rotated.astype(np.int32)
    os.system('''sed -e 's/ExonCount/spliced/g' %s > %s''' %(gem_path, gem_path_spliced))
    Spateo_refine_alignment(gem_path_spliced, img_path_rotated, rotated, out_dir, out_img_prefix)

def Spateo_refine_alignment(gem_path, img_path, rotated, out_dir, out_img_prefix):
    # create Spateo class
    adata = st.io.read_bgi_agg(
        gem_path, img_path,
        gene_agg={'nuclear': ['MALAT1', 'NETA1']},
        prealigned = False,
    )
    if 'spliced' in adata.layers:
        adata.layers['unspliced'] = adata.X - adata.layers['spliced']
    else:
        adata.layers['unspliced'] = adata.X
    fig, ax = plt.subplots(ncols=1, figsize=(20, 20), tight_layout=True)
    ax.imshow(adata.layers['stain'])
    st.pl.imshow(adata, 'unspliced', ax=ax, alpha=0.6, cmap='Reds', vmax=3, use_scale=False, save_show_or_return='return')
    ax.set_title('Before alignment', fontsize = 20)
    plt.savefig(os.path.join(out_dir, "%s_before_alignment.png" %out_img_prefix), bbox_inches='tight')
    # refine alignment
    before = adata.layers['stain'].copy()
    st.cs.refine_alignment(adata, mode='rigid', transform_layers=['stain'])
    fig, axes = plt.subplots(ncols=2, figsize=(20, 12), tight_layout=True)
    axes[0].imshow(before)
    st.pl.imshow(adata, 'unspliced', ax=axes[0], alpha=0.6, cmap='Reds', vmax=5, use_scale=False, save_show_or_return='return')
    axes[0].set_title('Before alignment', fontsize = 20)
    st.pl.imshow(adata, 'stain', ax=axes[1], use_scale=False, save_show_or_return='return')
    st.pl.imshow(adata, 'unspliced', ax=axes[1], alpha=0.6, cmap='Reds', vmax=5, use_scale=False, save_show_or_return='return')
    axes[1].set_title('After alignment', fontsize = 20)
    plt.savefig(os.path.join(out_dir, "%s_alignment_compare.png" %out_img_prefix), bbox_inches='tight')
    img_trans = np.pad(adata.layers['stain'], ((rotated.shape[0] - adata.layers['stain'].shape[0],0), (rotated.shape[1] - adata.layers['stain'].shape[1],0)))
    imsave(os.path.join(out_dir, "%s_regist_transposed_aligned_by_Spateo.tiff" %out_img_prefix), img_trans)

if __name__ == '__main__':
    parser = AlignmentParser()
    gem_path = parser.spot_expr_file
    img_path = parser.regist_tif_file
    num_workers = parser.num_workers
    out_dir = parser.out_dir
    out_prefix = parser.out_prefix
    # set number of threads
    Alignment(gem_path, img_path, num_workers, out_dir, out_prefix)

