# Simple Chromnitron Dataset - Just Two Matrices
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import List, Optional, Tuple

class ChromnitronDataset:
    """
    Simple dataset for two condition comparison
    User provides control and experiment matrices (loci x CAPs)
    """
    def __init__(self, control_matrix: np.ndarray, experiment_matrix: np.ndarray, 
                 gene_names: List[str], cap_names: List[str], 
                 control_name: str = 'Control', experiment_name: str = 'Experiment'):
        """
        Create ChromnitronDataset from two matrices
        
        Args:
            control_matrix: Control condition matrix (n_loci x n_caps)
            experiment_matrix: Experiment condition matrix (n_loci x n_caps) 
            gene_names: List of gene/loci names
            cap_names: List of CAP names
            control_name: Name for control condition
            experiment_name: Name for experiment condition
        """
        # Validate inputs
        if control_matrix.shape != experiment_matrix.shape:
            raise ValueError(f"Matrix shapes don't match: {control_matrix.shape} vs {experiment_matrix.shape}")
        
        if len(gene_names) != control_matrix.shape[0]:
            raise ValueError(f"Gene names length {len(gene_names)} doesn't match matrix rows {control_matrix.shape[0]}")
        
        if len(cap_names) != control_matrix.shape[1]:
            raise ValueError(f"CAP names length {len(cap_names)} doesn't match matrix cols {control_matrix.shape[1]}")
        
        # Store data
        self.matrices = {
            control_name: control_matrix.copy(),
            experiment_name: experiment_matrix.copy(),
            'differential': None  # Will be computed later
        }
        
        self.gene_names = gene_names.copy()
        self.cap_names = cap_names.copy()
        self.control_name = control_name
        self.experiment_name = experiment_name
        
        # Store filtered feature names
        self.filtered_gene_names = None
        self.filtered_cap_names = None
        
        # Metadata
        self.shape = control_matrix.shape
        self.n_genes, self.n_caps = self.shape
    
    def __repr__(self):
        return f"ChromnitronDataset(shape={self.shape}, conditions=['{self.control_name}', '{self.experiment_name}'])"
    
    # ==================== Basic Access ====================
    
    def get_matrix(self, condition: str) -> np.ndarray:
        """Get matrix for a condition"""
        if condition not in self.matrices:
            available = list(self.matrices.keys())
            raise ValueError(f"Condition '{condition}' not found. Available: {available}")
        if self.matrices[condition] is None:
            raise ValueError(f"Matrix for '{condition}' is None. Compute it first.")
        return self.matrices[condition]
    
    def get_control(self) -> np.ndarray:
        """Get control matrix"""
        return self.get_matrix(self.control_name)
    
    def get_experiment(self) -> np.ndarray:
        """Get experiment matrix"""
        return self.get_matrix(self.experiment_name)
    
    def get_differential(self) -> np.ndarray:
        """Get differential matrix"""
        return self.get_matrix('differential')
    
    def to_dataframe(self, condition: str) -> pd.DataFrame:
        """Convert matrix to DataFrame with proper labels"""
        matrix = self.get_matrix(condition)
        
        # Use filtered names if available
        if self.filtered_gene_names is not None and self.filtered_cap_names is not None:
            gene_names = self.filtered_gene_names
            cap_names = self.filtered_cap_names
        else:
            gene_names = self.gene_names
            cap_names = self.cap_names
            
        return pd.DataFrame(matrix, index=gene_names, columns=cap_names)
    
    # ==================== Operations ====================
    
    def log_transform(self, pseudocount: float = 0.01, conditions: List[str] = None, base: str = '2'):
        """Apply log transformation to specified conditions
        
        Args:
            pseudocount: Small value to add before log transform to avoid log(0)
            conditions: List of conditions to transform (default: both control and experiment)
            base: Logarithm base - 'e' for natural log, '2' for log2, '10' for log10
        """
        if conditions is None:
            conditions = [self.control_name, self.experiment_name]
        
        for condition in conditions:
            if condition in self.matrices and self.matrices[condition] is not None:
                if base == 'e':
                    self.matrices[condition] = np.log(self.matrices[condition] + pseudocount)
                elif base == '2':
                    self.matrices[condition] = np.log2(self.matrices[condition] + pseudocount)
                elif base == '10':
                    self.matrices[condition] = np.log10(self.matrices[condition] + pseudocount)
                else:
                    raise ValueError("Base must be 'e', '2', or '10'")
        return self
    
    def compute_differential(self, log_transform: bool = True, pseudocount: float = 0.01, base: str = '2'):
        """Compute differential: experiment - control
        
        Args:
            log_transform: If True, apply log transform before computing difference (recommended)
            pseudocount: Small value to add before log transform to avoid log(0)
            base: Logarithm base - 'e' for natural log, '2' for log2, '10' for log10
        """
        if log_transform:
            # Apply log transform without modifying original matrices
            if base == 'e':
                control_log = np.log(self.get_control() + pseudocount)
                experiment_log = np.log(self.get_experiment() + pseudocount)
            elif base == '2':
                control_log = np.log2(self.get_control() + pseudocount)
                experiment_log = np.log2(self.get_experiment() + pseudocount)
            elif base == '10':
                control_log = np.log10(self.get_control() + pseudocount)
                experiment_log = np.log10(self.get_experiment() + pseudocount)
            else:
                raise ValueError("Base must be 'e', '2', or '10'")
            
            # Compute log fold change
            self.matrices['differential'] = experiment_log - control_log
        else:
            # Simple difference without log transform
            self.matrices['differential'] = self.get_experiment() - self.get_control()
        
        return self
    
    def get_most_variable(self, n_genes: int, n_caps: int, condition: str = 'differential'):
        """Find most variable genes and CAPs based on specified condition (doesn't filter matrices)
        
        Args:
            n_genes: Number of most variable genes to identify
            n_caps: Number of most variable CAPs to identify
            condition: Which matrix to use for variance calculation 
                      ('differential', control_name, experiment_name, or 'average')
        
        Returns:
            Dict with gene_indices, cap_indices, gene_names, cap_names, and scores
        """
        if condition == 'differential':
            if self.matrices['differential'] is None:
                self.compute_differential()
            data = self.matrices['differential']
        elif condition == 'average':
            # Use average of both conditions
            data = (self.get_control() + self.get_experiment()) / 2
        elif condition in self.matrices and self.matrices[condition] is not None:
            data = self.matrices[condition]
        else:
            available = [k for k, v in self.matrices.items() if v is not None] + ['average']
            raise ValueError(f"Condition '{condition}' not available. Available: {available}")
        
        # Calculate variance
        gene_var = np.var(data, axis=1)  # Variance across CAPs for each gene
        cap_var = np.var(data, axis=0)   # Variance across genes for each CAP
        
        # Get top variable features
        top_gene_indices = np.argsort(gene_var)[-n_genes:]
        top_cap_indices = np.argsort(cap_var)[-n_caps:]
        
        # Get names for the selected features
        selected_gene_names = [self.gene_names[i] for i in top_gene_indices]
        selected_cap_names = [self.cap_names[i] for i in top_cap_indices]
        
        return {
            'gene_indices': top_gene_indices,
            'cap_indices': top_cap_indices,
            'gene_names': selected_gene_names,
            'cap_names': selected_cap_names,
            'gene_scores': gene_var[top_gene_indices],
            'cap_scores': cap_var[top_cap_indices],
            'based_on': condition,
            'n_genes': n_genes,
            'n_caps': n_caps
        }
    
    def apply_feature_filter(self, gene_indices: np.ndarray, cap_indices: np.ndarray, 
                           gene_names: List[str] = None, cap_names: List[str] = None):
        """Apply feature filtering to all matrices using provided indices
        
        Args:
            gene_indices: Array of gene indices to keep
            cap_indices: Array of CAP indices to keep  
            gene_names: Optional list of gene names (will be inferred if not provided)
            cap_names: Optional list of CAP names (will be inferred if not provided)
        """
        # Filter all matrices
        for matrix_name, matrix in self.matrices.items():
            if matrix is not None:
                self.matrices[matrix_name] = matrix[np.ix_(gene_indices, cap_indices)]
        
        # Update feature names
        if gene_names is not None and cap_names is not None:
            self.filtered_gene_names = gene_names
            self.filtered_cap_names = cap_names
        else:
            self.filtered_gene_names = [self.gene_names[i] for i in gene_indices]
            self.filtered_cap_names = [self.cap_names[i] for i in cap_indices]
        
        # Update shape
        self.shape = self.matrices[self.control_name].shape
        self.n_genes, self.n_caps = self.shape
        
        return self
    
    def variance_filter(self, n_genes: int, n_caps: int, condition: str = 'differential'):
        """Find most variable features and apply filtering (convenience function)
        
        This is equivalent to: get_most_variable() followed by apply_feature_filter()
        """
        var_results = self.get_most_variable(n_genes, n_caps, condition)
        self.apply_feature_filter(
            var_results['gene_indices'], 
            var_results['cap_indices'],
            var_results['gene_names'],
            var_results['cap_names']
        )
        
        return var_results
    
    def remove_noise(self, top_percent: float = 10.0, conditions: List[str] = None):
        """Remove noise by keeping only extreme values"""
        if conditions is None:
            conditions = ['differential']  # Default to differential only
        elif conditions == ['all']:
            # Special case: apply to all available matrices
            conditions = [k for k, v in self.matrices.items() if v is not None]
        
        for condition in conditions:
            if condition not in self.matrices or self.matrices[condition] is None:
                print(f"Warning: Condition '{condition}' not available, skipping")
                continue
            
            data = self.matrices[condition].copy()
            
            # For differential data, keep top and bottom percentiles
            if condition == 'differential':
                # Per column (CAP)
                for i in range(data.shape[1]):
                    col_data = data[:, i]
                    percentile_bot = np.percentile(col_data, top_percent)
                    percentile_top = np.percentile(col_data, 100 - top_percent)
                    mask = np.logical_and(col_data >= percentile_bot, col_data < percentile_top)
                    data[mask, i] = 0
                
                # Per row (gene)
                for i in range(data.shape[0]):
                    row_data = data[i, :]
                    percentile_bot = np.percentile(row_data, top_percent)
                    percentile_top = np.percentile(row_data, 100 - top_percent)
                    mask = np.logical_and(row_data >= percentile_bot, row_data < percentile_top)
                    data[i, mask] = 0
            
            else:
                # For individual conditions, keep only top percentiles
                # Per column (CAP)
                for i in range(data.shape[1]):
                    col_data = data[:, i]
                    percentile_top = np.percentile(col_data, 100 - top_percent)
                    data[col_data < percentile_top, i] = 0
                
                # Per row (gene)
                for i in range(data.shape[0]):
                    row_data = data[i, :]
                    percentile_top = np.percentile(row_data, 100 - top_percent)
                    data[i, row_data < percentile_top] = 0
            
            self.matrices[condition] = data
        
        return self
    
    # ==================== Analysis Methods ====================
    
    def get_signature_data(self, method: str = 'kmeans'):
        """Get appropriate data for signature analysis"""
        if method in ['kmeans', 'svd']:
            # Use differential matrix for clustering methods
            if self.matrices['differential'] is None:
                self.compute_differential()
            matrix = self.get_differential()
            gene_names = self.filtered_gene_names or self.gene_names
            cap_names = self.filtered_cap_names or self.cap_names
            return matrix, gene_names, cap_names
            
        elif method == 'nmf':
            # Use individual matrices for NMF
            control_matrix = self.get_control()
            experiment_matrix = self.get_experiment()
            gene_names = self.filtered_gene_names or self.gene_names
            cap_names = self.filtered_cap_names or self.cap_names
            
            return {
                self.control_name: (control_matrix, gene_names, cap_names),
                self.experiment_name: (experiment_matrix, gene_names, cap_names)
            }
        
        else:
            raise ValueError(f"Unknown method: {method}")
    
    # ==================== I/O Methods ====================
    
    def save(self, output_path: str):
        """Save dataset"""
        os.makedirs(output_path, exist_ok=True)
        
        # Save all matrices
        for condition, matrix in self.matrices.items():
            if matrix is not None:
                # Numpy format
                np.save(os.path.join(output_path, f'{condition}.npy'), matrix)
                
                # CSV format
                df = self.to_dataframe(condition)
                df.to_csv(os.path.join(output_path, f'{condition}.csv'))
        
        # Save metadata as json
        metadata = {
            'gene_names': self.gene_names,
            'cap_names': self.cap_names,
            'filtered_gene_names': self.filtered_gene_names,
            'filtered_cap_names': self.filtered_cap_names,
            'control_name': self.control_name,
            'experiment_name': self.experiment_name,
            'shape': self.shape
        }
        
        import json
        with open(os.path.join(output_path, 'metadata.json'), 'w') as f:
            json.dump(metadata, f, indent=2, default=str)
    
    @classmethod
    def load(cls, input_path: str):
        """Load dataset from saved files"""
        # Load metadata
        import json
        with open(os.path.join(input_path, 'metadata.json'), 'r') as f:
            metadata = json.load(f)
        
        # Load matrices
        control_matrix = np.load(os.path.join(input_path, f"{metadata['control_name']}.npy"))
        experiment_matrix = np.load(os.path.join(input_path, f"{metadata['experiment_name']}.npy"))
        
        # Create instance
        instance = cls(
            control_matrix, experiment_matrix,
            metadata['gene_names'], metadata['cap_names'],
            metadata['control_name'], metadata['experiment_name']
        )
        
        # Restore filtered names
        instance.filtered_gene_names = metadata.get('filtered_gene_names')
        instance.filtered_cap_names = metadata.get('filtered_cap_names')
        
        # Load differential if it exists
        diff_file = os.path.join(input_path, 'differential.npy')
        if os.path.exists(diff_file):
            instance.matrices['differential'] = np.load(diff_file)
        
        return instance
    
    # ==================== Convenience Methods ====================
    
    def summary(self):
        """Print dataset summary"""
        print(f"ChromnitronDataset Summary:")
        print(f"  Shape: {self.shape} ({self.n_genes} genes × {self.n_caps} CAPs)")
        print(f"  Conditions: {self.control_name}, {self.experiment_name}")
        
        if self.filtered_gene_names is not None:
            print(f"  Filtered: {len(self.filtered_gene_names)} genes × {len(self.filtered_cap_names)} CAPs")
        
        available_matrices = [k for k, v in self.matrices.items() if v is not None]
        print(f"  Available matrices: {available_matrices}")
        
        for condition, matrix in self.matrices.items():
            if matrix is not None:
                print(f"    {condition}: range [{np.min(matrix):.3f}, {np.max(matrix):.3f}]")


# ==================== Factory Functions ====================

def create_from_files(control_file: str, experiment_file: str, 
                     gene_names_file: str, cap_names_file: str,
                     control_name: str = 'Control', experiment_name: str = 'Experiment'):
    """Create ChromnitronDataset from files"""
    
    # Load matrices
    control_matrix = np.load(control_file)
    experiment_matrix = np.load(experiment_file)
    
    # Load gene names
    if gene_names_file.endswith('.csv'):
        gene_df = pd.read_csv(gene_names_file)
        if 'transcript_name' in gene_df.columns:
            gene_names = gene_df['transcript_name'].tolist()
        else:
            gene_names = gene_df.iloc[:, 0].tolist()
    else:
        with open(gene_names_file, 'r') as f:
            gene_names = f.read().strip().split('\n')
    
    # Load CAP names
    if cap_names_file.endswith('.csv'):
        cap_names = pd.read_csv(cap_names_file, header=None)[0].tolist()
    else:
        with open(cap_names_file, 'r') as f:
            cap_names = f.read().strip().split('\n')
    
    return ChromnitronDataset(control_matrix, experiment_matrix, gene_names, cap_names,
                             control_name, experiment_name)


def create_from_original_data(data_file: str, gene_file: str, cap_file: str, celltype_file: str,
                             mode: str = 'max'):
    """Create ChromnitronDataset from original 4D data structure"""
    
    # Load data
    raw_data = np.load(data_file)  # Shape: (n_celltypes, n_loci, n_caps, 2)
    raw_data = raw_data.reshape(2, -1, *raw_data.shape[1:]).mean(axis=1)   #hua added
    # Load metadata
    gene_df = pd.read_csv(gene_file)
    gene_names = gene_df['transcript_name'].tolist()
    
    cap_names = pd.read_csv(cap_file, header=None)[0].tolist()
    
    with open(celltype_file, 'r') as f:
        celltype_names = f.read().strip().split('\n')
    
    # Extract matrices for two cell types
    if len(celltype_names) < 2:
        raise ValueError("Need at least 2 cell types")
    
    # Select max or mean values
    if mode == 'max':
        control_matrix = raw_data[1, :, :, 0]     # Second celltype, max values
        experiment_matrix = raw_data[0, :, :, 0]  # First celltype, max values
    elif mode == 'mean':
        control_matrix = raw_data[1, :, :, 1]     # Second celltype, mean values
        experiment_matrix = raw_data[0, :, :, 1]  # First celltype, mean values
    else:
        raise ValueError("Mode must be 'max' or 'mean'")
    
    # Create dataset
    config_name = data_file.split('/')[-2]
    if '_vs_' in config_name:
        control_name = config_name.split('_vs_')[0]
        experiment_name = config_name.split('_vs_')[1]
    else:
        control_name = config_name + '_rep1'
        experiment_name = config_name + '_rep2'
    dataset = ChromnitronDataset(control_matrix, experiment_matrix, gene_names, cap_names,
                                 control_name, experiment_name) #hua changed
                                #celltype_names[1], celltype_names[0])
    
    return dataset


# ==================== Usage Examples ====================

def example_usage():
    """Example usage of the simple dataset"""
    
    # Option 1: Create from prepared matrices
    control_matrix = np.random.rand(1000, 500)
    experiment_matrix = np.random.rand(1000, 500)
    gene_names = [f"Gene_{i}" for i in range(1000)]
    cap_names = [f"CAP_{i}" for i in range(500)]
    
    dataset = ChromnitronDataset(control_matrix, experiment_matrix, 
                                gene_names, cap_names, 'Tex', 'Teff')
    
    # Option 2: Create from files
    # dataset = create_from_files('control.npy', 'experiment.npy', 'genes.txt', 'caps.txt')
    
    # Option 3: Create from original data structure (no auto log transform now)
    # dataset = create_from_original_data('upregulated_data.npy', 'genes.csv', 'caps.txt', 'celltypes.txt', log_transform=False)
    
    # Compute log2 fold change differential (individual matrices stay raw for NMF)
    dataset.compute_differential(log_transform=True, base='2')
    
    # Option A: Simple variance filtering (same as before)
    dataset.variance_filter(n_genes=500, n_caps=100, condition='differential')
    
    # Option B: Advanced variance filtering workflow
    # Find different sets of variable features
    # diff_vars = dataset.get_most_variable(n_genes=500, n_caps=100, condition='differential')
    # tex_vars = dataset.get_most_variable(n_genes=300, n_caps=80, condition='Tex')
    # avg_vars = dataset.get_most_variable(n_genes=800, n_caps=150, condition='average')
    
    # Compare or combine different strategies
    # common_genes = np.intersect1d(diff_vars['gene_indices'], tex_vars['gene_indices'])
    # print(f"Genes variable in both differential and Tex: {len(common_genes)}")
    
    # Apply chosen filtering
    # dataset.apply_feature_filter(diff_vars['gene_indices'], diff_vars['cap_indices'])
    
    # Remove noise from specific or all matrices
    dataset.remove_noise(top_percent=10.0, conditions=['all'])  # Apply to all matrices
    # dataset.remove_noise(top_percent=5.0, conditions=['differential'])  # Only differential
    
    # Get data for different signature analysis methods
    
    # For clustering methods (uses log2fc differential matrix)
    diff_matrix, genes, caps = dataset.get_signature_data('kmeans')
    print(f"Differential matrix for clustering: {diff_matrix.shape}")
    print(f"Log2FC range: [{np.min(diff_matrix):.2f}, {np.max(diff_matrix):.2f}]")
    
    # For NMF (uses raw individual matrices - non-negative)
    nmf_data = dataset.get_signature_data('nmf')
    print(f"NMF data keys: {list(nmf_data.keys())}")
    for condition, (matrix, g, c) in nmf_data.items():
        print(f"  {condition}: {matrix.shape}, range: [{np.min(matrix):.2f}, {np.max(matrix):.2f}]")
    
    # Save results
    output_path = '/path/to/output'
    dataset.save(output_path)
    print(f"Saved dataset to {output_path}")
    
    # Load later
    # loaded_dataset = ChromnitronDataset.load(output_path)
    
    # Summary
    dataset.summary()
    
    # Example of accessing specific matrices
    print(f"\nMatrix access examples:")
    print(f"Control (raw): {dataset.get_control().shape}")
    print(f"Experiment (raw): {dataset.get_experiment().shape}")  
    print(f"Differential (log2fc): {dataset.get_differential().shape}")
    
    # Convert to DataFrame for analysis
    diff_df = dataset.to_dataframe('differential')
    print(f"Differential DataFrame: {diff_df.shape}")
    print(f"Gene names: {diff_df.index[:5].tolist()}")
    print(f"CAP names: {diff_df.columns[:5].tolist()}")

if __name__ == "__main__":
    print("Simple ChromnitronDataset")
    print("See example_usage() for usage patterns")