import torch, numpy, scanpy, pandas
from tqdm import tqdm
from sklearn.cluster import KMeans
from .positive_pair import build_pair
from .negative_sample import nega_sample
from .utils import contrastive_loss

class train():
    def __init__(self,
                 processed_data: numpy.ndarray,
                 graph_dict: dict,
                 model: torch.nn.Module,
                 pre_epochs: int = 200,
                 epochs: int = 200,
                 corrupt: float = 0.001,
                 lr: float = 5e-4,
                 weight_decay: float = 1e-4,
                 domains: numpy.ndarray = None,
                 kl_weight: float = 100,
                 mse_weight: float = 10,
                 bce_kld_weight: float = 0.1,
                 domain_weight: float = 1,
                 use_gpu: bool = True,
                 use_contra: bool = False
                ):
        """
        :param processed_data: numpy.ndarray, Processed input data [n_spots, n_features]
        :param graph_dict: dict, Dictionary containing graph information with keys:
            - 'adj_norm': Normalized adjacency matrix
            - 'adj_label': Original adjacency labels
            - 'norm_value': Normalization factor
        :param model: torch.nn.Module, stKAN model instance
        :param pre_epochs: int, Number of pretraining epochs. default: int=200.
        :param epochs: int, Number of training epochs. default: int=200.
        :param corrupt: float, Corruption rate for masking noise. default: float=0.001.
        :param lr: float, Learning rate. default: float=5e-4.
        :param weight_decay: float, Weight decay for optimizer. default: float=1e-4.
        :param domains: numpy.ndarray, optional, Domain labels for adversarial training. default: None.
        :param kl_weight: float, Weight for KL divergence loss. default: int=100.
        :param mse_weight: float, Weight for MSE reconstruction loss. default: int=10.
        :param bce_kld_weight: float, Weight for BCE and KLD losses. default: float=0.1.
        :param domain_weight: float, Weight for domain classification loss. default: int=1.
        :param use_gpu: bool, Whether to use GPU if available. default: bool=True.
        :param use_contra: bool, Whether to use contrastive learning. default: bool=False.
        """
        # Device setup
        self.device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
        
        # Data setup
        self.processed_data = processed_data
        self.data = torch.FloatTensor(processed_data).to(self.device)
        self.adj = graph_dict['adj_norm'].to(self.device)
        self.adj_label = graph_dict['adj_label'].to(self.device)
        self.norm = graph_dict['norm_value']
        self.num_spots = self.data.shape[0]
        
        # Model setup
        self.model = model.to(self.device)
        self.optimizer = torch.optim.Adam(
            self.model.parameters(), 
            lr=lr, 
            weight_decay=weight_decay
        )
        
        # Training parameters
        self.pre_epochs = pre_epochs
        self.epochs = epochs
        self.corrupt = corrupt
        self.kl_weight = kl_weight
        self.mse_weight = mse_weight
        self.bce_kld_weight = bce_kld_weight
        self.domain_weight = domain_weight

        # Contrastive learning
        self.use_contra = use_contra
        self.knn = 25
        self.num_centroids = 20
        self.clus_num_iters = 20
        self.pool_percent = 0.95
        self.sample_percent = 0.1
        self.topkdevice = self.device
        self.temp = 0.5
        self.alpha = 0.2
        
        # Clustering parameters
        self.dec_tol = 0.001  # Tolerance for cluster assignment change
        self.q_stride = 20  # Interval for updating target distribution
        
        # Domain adaptation
        self.domains = None
        if domains is not None:
            self.domains = torch.from_numpy(domains).to(self.device)

    def pretrain(self, grad_clip: float = 5.0) -> None:
        """
        :param grad_clip: float, Maximum gradient norm for clipping. default: float=5.0.
        """
        with tqdm(total=self.pre_epochs, 
                 desc="Pretraining initial model",
                 bar_format="{l_bar}{bar} [ time left: {remaining} ]", ncols=80, dynamic_ncols=True, leave=True) as pbar:
            
            for epoch in range(self.pre_epochs):
                # Add masking noise for robustness
                inputs_corr = self._masking_noise(self.data, self.corrupt)
                
                # Forward pass
                self.model.train()
                self.optimizer.zero_grad()
                
                if self.domains is None:
                    z, mu, logvar, de_feat, _, _, _ = self.model(inputs_corr, self.adj)
                    preds = self.model.dc(z)
                else:
                    z, mu, logvar, de_feat, _, _, _, domain_pred = self.model(inputs_corr, self.adj)
                    preds = self.model.model.dc(z)
                
                if self.use_contra:
                    emb = z.clone()
                    with torch.no_grad():
                        positive = build_pair(self.device, self.num_centroids, self.clus_num_iters)
                        positive_emb, similarity, knn_neighbor = positive.fit(emb, self.knn)
                        nega_sampler = nega_sample(knn_neighbor, self.device, self.pool_percent, self.sample_percent, self.topkdevice)
                        nega_sampler.fit(similarity, 0)
                        loss_contra = self.alpha * contrastive_loss(z, positive_emb, nega_sampler.negasample, nega_sampler.n_nega, self.device, temperature=self.temp)
                
                # Compute loss
                loss = self.model.stKAN_loss(
                            decoded=de_feat,
                            x=self.data,
                            preds=preds,
                            labels=self.adj_label,
                            mu=mu,
                            logvar=logvar,
                            n_nodes=self.num_spots,
                            norm=self.norm,
                            mask=self.adj_label,
                            mse_weight=self.mse_weight,
                            bce_kld_weight=self.bce_kld_weight
                            )
                
                # Add domain loss if applicable
                if self.domains is not None:
                    domain_loss = torch.nn.functional.cross_entropy(domain_pred, self.domains)
                    loss += self.domain_weight * domain_loss
                    if self.use_contra:
                        loss += loss_contra
                
                # Backward pass
                if self.use_contra:
                    loss += loss_contra
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), grad_clip)
                self.optimizer.step()
                
                pbar.update(1)

    @torch.no_grad()
    def process(self) -> tuple:
        """
        :return: tuple: (z, q)
            z: numpy.ndarray, Latent representation [n_spots, n_features].
            q: numpy.ndarray, Soft cluster assignments [n_spots, n_clusters].
        """
        self.model.eval()
        if self.domains is None:
            z, _, _, _, q, _, _ = self.model(self.data, self.adj)
        else:
            z, _, _, _, q, _, _, _ = self.model(self.data, self.adj)
        return z.cpu().numpy(), q.cpu().numpy()

    def save_model(self, save_path: str) -> None:
        torch.save({'state_dict': self.model.state_dict()}, save_path)
        print(f'Saved model to {save_path}')

    def load_model(self, load_path: str) -> None:
        state_dict = torch.load(load_path)
        self.model.load_state_dict(state_dict['state_dict'])
        print(f'Loaded model from {load_path}')

    def fit(self, 
            cluster_n: int = 20,
            cluster_type: str = 'Louvain',
            resolution: float = 1.0,
            pretrain: bool = True) -> numpy.ndarray:
        """
        :param cluster_n: int, Number of clusters. default: int=20.
        :param cluster_type : str, Clustering method ['Louvain', 'KMeans']. default: str='Louvain'.
        :param resolution: float, Resolution parameter for Louvain clustering. default: float=1.0.
        :param pretrain: bool, Whether to run pretraining. default: bool=True.
            
        :return: numpy.ndarray, Final cluster assignments [n_spots].
        """
        # Pretrain if specified
        if pretrain:
            self.pretrain()
            pre_z, _ = self.process()
            
        # Initialize cluster centers
        y_pred_last = self._initialize_clusters(pre_z, cluster_n, cluster_type, resolution)
        
        # Main training loop
        with tqdm(total=self.epochs,
                  desc="Training final model",
                  bar_format="{l_bar}{bar} [ time left: {remaining} ]", ncols=80, dynamic_ncols=True, leave=True) as pbar:
            
            for epoch in range(self.epochs):
                # Update target distribution periodically
                if epoch % self.q_stride == 0:
                    _, q = self.process()
                    q = self.model.target_distribution(torch.Tensor(q).to(self.device))
                    y_pred = q.cpu().numpy().argmax(1)
                    
                    # Check for convergence
                    delta_label = numpy.sum(y_pred != y_pred_last) / y_pred.shape[0]
                    y_pred_last = y_pred.copy()
                    if epoch > 0 and delta_label < self.dec_tol:
                        print(f'delta_label {delta_label:.4f} < tol {self.dec_tol}')
                        print('Reached tolerance threshold. Stopping training.')
                        break
                
                # Training step
                self._train_step(q)
                pbar.update(1)
        
        # Return final cluster assignments
        _, q = self.process()
        return q.argmax(1)

    def _train_step(self, q: torch.Tensor) -> None:
        """Perform a single training step"""
        self.model.train()
        self.optimizer.zero_grad()
        
        # Forward pass
        if self.domains is None:
            z, mu, logvar, de_feat, out_q, _, _ = self.model(self.data, self.adj)
            preds = self.model.dc(z)
        else:
            z, mu, logvar, de_feat, out_q, _, _, domain_pred = self.model(self.data, self.adj)
            preds = self.model.model.dc(z)
        
        if self.use_contra:
            emb = z.clone()
            with torch.no_grad():
                positive = build_pair(self.device, self.num_centroids, self.clus_num_iters)
                positive_emb, similarity, knn_neighbor = positive.fit(emb, self.knn)
                nega_sampler = nega_sample(knn_neighbor, self.device, self.pool_percent, self.sample_percent, self.topkdevice)
                nega_sampler.fit(similarity, 0)
                loss_contra = self.alpha * contrastive_loss(z, positive_emb, nega_sampler.negasample, nega_sampler.n_nega, self.device, temperature=self.temp)
            
        # Compute losses
        loss_stKAN = self.model.stKAN_loss(
            decoded = de_feat,
            x = self.data,
            preds = preds,
            labels = self.adj_label,
            mu = mu,
            logvar = logvar,
            n_nodes = self.num_spots,
            norm = self.norm,
            mask = self.adj_label,
            mse_weight = self.mse_weight,
            bce_kld_weight = self.bce_kld_weight
        )
        
        loss_kl = torch.nn.functional.kl_div(out_q.log(), q, reduction='batchmean')
        
        # Combine losses
        if self.domains is None:
            loss = self.kl_weight * loss_kl + loss_stKAN
            if self.use_contra:
                loss += loss_contra
        else:
            domain_loss = torch.nn.functional.cross_entropy(domain_pred, self.domains)
            loss = self.kl_weight * loss_kl + loss_stKAN + self.domain_weight * domain_loss
            if self.use_contra:
                loss += loss_contra
        
        # Backward pass
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5)
        self.optimizer.step()

    def _initialize_clusters(self, 
                           z: numpy.ndarray,
                           cluster_n: int,
                           cluster_type: str,
                           resolution: float) -> numpy.ndarray:
        if cluster_type == 'KMeans':
            kmeans = KMeans(n_clusters=cluster_n, n_init=cluster_n*2, random_state=88)
            y_pred = kmeans.fit_predict(z)
            centers = kmeans.cluster_centers_
        elif cluster_type == 'Louvain':
            cluster_data = scanpy.AnnData(z)
            scanpy.pp.neighbors(cluster_data, n_neighbors=cluster_n)
            scanpy.tl.louvain(cluster_data, resolution=resolution)
            y_pred = cluster_data.obs['louvain'].astype(int).to_numpy()
            
            # Calculate cluster centers
            features = pandas.DataFrame(z, index=numpy.arange(z.shape[0]))
            groups = pandas.Series(y_pred, index=numpy.arange(z.shape[0]), name="Group")
            merged = pandas.concat([features, groups], axis=1)
            centers = numpy.asarray(merged.groupby("Group").mean())
        
        # Update cluster layer
        target_model = self.model if self.domains is None else self.model.model
        target_model.cluster_layer.data = torch.tensor(centers).to(self.device)
        return y_pred

    @staticmethod
    def _masking_noise(data: torch.Tensor, frac: float) -> torch.Tensor:
        data_noise = data.clone()
        mask = torch.rand(data.size()) < frac
        data_noise[mask.to(data.device)] = 0
        return data_noise