import pandas as pd
import numpy as np
from collections import Counter
import json
from typing import Dict, List, Union, Any
from pathlib import Path
from transformers import AutoTokenizer
from tqdm import tqdm

class CorpusTokenAnalyzer:
    def __init__(self, model_name_or_tokenizer: Union[str, Any], corpus_data: List[Dict] = None, corpus_file: str = None, output_corpus_file: str = None):
        """
        Initialize the CorpusTokenAnalyzer.
        
        Args:
            model_name_or_tokenizer: Model name (str) for AutoTokenizer or tokenizer object
            corpus_data (List[Dict]): List of documents with 'chunk_id', 'text', and 'title' fields
            corpus_file (str): Path to the corpus file (CSV, JSON, or JSONL)
            output_corpus_file (str): Path for the output corpus file with context lengths added
        """
        self.corpus_data = corpus_data
        self.corpus_file = corpus_file
        self.output_corpus_file = output_corpus_file or self._generate_output_filename()
        self.chunks_data = []
        self.chunk_context_lengths = {}
        self.token_doc_matrix = None
        self.chunk_ids = []
        
        # Initialize tokenizer
        if isinstance(model_name_or_tokenizer, str):
            self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_tokenizer)
            self.model_name = model_name_or_tokenizer
        else:
            self.tokenizer = model_name_or_tokenizer
            self.model_name = getattr(model_name_or_tokenizer, 'name_or_path', 'unknown_model')
        
        # Get model vocabulary size
        self.vocab_size = self.tokenizer.vocab_size
        print(f"Using model: {self.model_name}")
        print(f"Model vocabulary size: {self.vocab_size}")
        
        if not self.corpus_data:
        # Load the corpus
            assert self.corpus_file is not None, "Corpus file must be provided if corpus_data is empty"
            self._load_corpus()
        else:
            self.chunks_data = self.corpus_data
        # Process tokens and build matrix
        self._process_tokens()
        # Add context lengths to original data and save
        self._add_context_lengths_to_corpus()
    
    def _generate_output_filename(self) -> str:
        """Generate output filename based on input filename."""
        file_path = Path(self.corpus_file)
        return str(file_path.parent / f"{file_path.stem}_with_context_length{file_path.suffix}")
    
    def _load_corpus(self):
        """Load corpus from file based on file extension."""
        file_path = Path(self.corpus_file)
        if file_path.suffix.lower() == '.csv':
            df = pd.read_csv(self.corpus_file)
            self.chunks_data = df.to_dict('records')
            self.file_format = 'csv'
        elif file_path.suffix.lower() == '.json':
            with open(self.corpus_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                if isinstance(data, list):
                    self.chunks_data = data
                    self.file_format = 'json'
                else:
                    raise ValueError("JSON file should contain a list of documents")
        elif file_path.suffix.lower() == '.jsonl':
            with open(self.corpus_file, 'r', encoding='utf-8') as f:
                self.chunks_data = [json.loads(line) for line in f]
                self.file_format = 'jsonl'
        else:
            raise ValueError("Unsupported file format. Use CSV, JSON, or JSONL")
        
        # Validate required columns
        required_fields = ['chunk_id', 'text', 'title']
        if not all(field in self.chunks_data[0] for field in required_fields):
            raise ValueError(f"Corpus must contain columns: {required_fields}")
    
    def _get_tokens(self, text: str) -> List[int]:
        """
        Get token IDs from text using the model tokenizer.
        
        Args:
            text (str): Input text
            
        Returns:
            List[int]: List of token IDs according to model vocabulary
        """
        # Use the transformers tokenizer encode method
        return self.tokenizer.encode(text, add_special_tokens=True)
    
    def _process_tokens(self):
        """Process all chunks and build token-document matrix using model vocabulary."""
        print("Processing tokens and building token-document matrix...")
        
        # Initialize token-document matrix: vocab_size x num_documents
        num_docs = len(self.chunks_data)
        self.token_doc_matrix = np.zeros((self.vocab_size, num_docs), dtype=np.int32)
        
        # Process each document
        for doc_idx, chunk in tqdm(enumerate(self.chunks_data), total=len(self.chunks_data), desc="Processing documents", unit="doc"):
            chunk_id = chunk['chunk_id']
            self.chunk_ids.append(chunk_id)
            
            # Combine title and contents
            full_text = f"{chunk['text']}"
            
            # Get token IDs using model tokenizer
            token_ids = self._get_tokens(full_text)
            
            # Store context length
            self.chunk_context_lengths[chunk_id] = len(token_ids)
            
            # Count tokens in this document
            token_counts = Counter(token_ids)
            
            # Fill the token-document matrix
            for token_id, count in token_counts.items():
                # Ensure token_id is within vocabulary range
                if 0 <= token_id < self.vocab_size:
                    self.token_doc_matrix[token_id, doc_idx] = count
                else:
                    print(f"Warning: Token ID {token_id} is outside vocabulary range")
        
        print(f"Token-document matrix shape: {self.token_doc_matrix.shape}")
        print(f"Matrix sparsity: {1 - np.count_nonzero(self.token_doc_matrix) / self.token_doc_matrix.size:.4f}")
    
    def _add_context_lengths_to_corpus(self):
        """Add context lengths to the original corpus data and save to file."""
        # Add context_length to each chunk
        for chunk in self.chunks_data:
            chunk_id = chunk['chunk_id']
            chunk['context_length'] = self.chunk_context_lengths[chunk_id]
        
        # Save based on original file format
        self._save_as_jsonl()

    def _save_as_jsonl(self):
        """Save the updated corpus as JSONL file."""
        with open(self.output_corpus_file, 'w', encoding='utf-8') as f:
            for chunk in self.chunks_data:
                f.write(json.dumps(chunk, ensure_ascii=False) + '\n')
        print(f"Updated corpus saved to: {self.output_corpus_file}")
    
    def _save_as_json(self):
        """Save the updated corpus as JSON file."""
        with open(self.output_corpus_file, 'w', encoding='utf-8') as f:
            json.dump(self.chunks_data, f, ensure_ascii=False, indent=2)
        print(f"Updated corpus saved to: {self.output_corpus_file}")
    
    def _save_as_csv(self):
        """Save the updated corpus as CSV file."""
        df = pd.DataFrame(self.chunks_data)
        df.to_csv(self.output_corpus_file, index=False)
        print(f"Updated corpus saved to: {self.output_corpus_file}")
    
    def get_context_lengths(self) -> Dict[str, int]:
        """Get context lengths for all chunks."""
        return self.chunk_context_lengths.copy()
    
    def get_token_doc_matrix(self) -> np.ndarray:
        """
        Get the token-document matrix.
        
        Returns:
            np.ndarray: Matrix where rows are token IDs (0 to vocab_size-1) 
                       and columns are documents, values are token counts
        """
        return self.token_doc_matrix.copy()
    
    def get_binary_token_doc_matrix(self) -> np.ndarray:
        """Get binary token-document matrix (1 if token appears in doc, 0 otherwise)."""
        return (self.token_doc_matrix > 0).astype(int)
    
    def get_chunk_token_map(self, chunk_id: str) -> Dict[int, int]:
        """
        Get token count map for a specific chunk.
        
        Args:
            chunk_id (str): ID of the chunk
            
        Returns:
            Dict[int, int]: Mapping of token_id to count in the document
        """
        if chunk_id not in self.chunk_ids:
            raise ValueError(f"Chunk ID '{chunk_id}' not found")
        
        doc_idx = self.chunk_ids.index(chunk_id)
        token_map = {}
        
        for token_id in range(self.vocab_size):
            count = self.token_doc_matrix[token_id, doc_idx]
            if count > 0:
                token_map[token_id] = count
        
        return token_map
    
    def get_token_info(self, token_id: int) -> Dict[str, Any]:
        """
        Get information about a specific token.
        
        Args:
            token_id (int): Token ID from model vocabulary
            
        Returns:
            Dict: Token information including text representation and document frequencies
        """
        if not 0 <= token_id < self.vocab_size:
            raise ValueError(f"Token ID {token_id} is outside vocabulary range [0, {self.vocab_size})")
        
        # Get token text representation
        try:
            token_text = self.tokenizer.decode([token_id])
        except:
            token_text = f"<token_{token_id}>"
        
        # Get document frequencies
        token_counts = self.token_doc_matrix[token_id, :]
        doc_frequency = np.sum(token_counts > 0)
        total_frequency = np.sum(token_counts)
        
        return {
            'token_id': token_id,
            'token_text': token_text,
            'document_frequency': int(doc_frequency),
            'total_frequency': int(total_frequency),
            'appears_in_docs': [self.chunk_ids[i] for i in range(len(self.chunk_ids)) 
                               if token_counts[i] > 0]
        }
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get overall statistics about the corpus and token analysis."""
        context_lengths = list(self.chunk_context_lengths.values())
        
        # Calculate additional statistics
        doc_frequencies = np.sum(self.token_doc_matrix > 0, axis=1)  # How many docs each token appears in
        total_frequencies = np.sum(self.token_doc_matrix, axis=1)    # Total count of each token
        
        return {
            'model_name': self.model_name,
            'model_vocab_size': self.vocab_size,
            'num_documents': len(self.chunks_data),
            'avg_context_length': float(np.mean(context_lengths)),
            'max_context_length': int(np.max(context_lengths)),
            'min_context_length': int(np.min(context_lengths)),
            'total_tokens': int(np.sum(context_lengths)),
            'unique_tokens_used': int(np.sum(doc_frequencies > 0)),  # How many vocab tokens actually appear
            'vocab_coverage': float(np.sum(doc_frequencies > 0) / self.vocab_size),  # % of vocab used
            'avg_tokens_per_doc': float(np.mean(np.sum(self.token_doc_matrix, axis=0))),
            'matrix_sparsity': float(1 - np.count_nonzero(self.token_doc_matrix) / self.token_doc_matrix.size),
            'most_frequent_token_id': int(np.argmax(total_frequencies)),
            'max_doc_frequency': int(np.max(doc_frequencies)),
            'output_file': self.output_corpus_file
        }
    
    def save_token_doc_matrix(self, output_file: str, format: str = 'numpy'):
        """
        Save the token-document matrix to file.
        
        Args:
            output_file (str): Output file path
            format (str): Format to save ('numpy', 'csv', 'sparse')
        """
        if format == 'numpy':
            np.save(output_file, self.token_doc_matrix)
        elif format == 'csv':
            # Create DataFrame with token IDs as index and chunk IDs as columns
            df = pd.DataFrame(
                self.token_doc_matrix,
                index=[f"token_{i}" for i in range(self.vocab_size)],
                columns=self.chunk_ids
            )
            df.to_csv(output_file)
        elif format == 'sparse':
            # Save as sparse matrix (COO format)
            from scipy import sparse
            sparse_matrix = sparse.coo_matrix(self.token_doc_matrix)
            sparse.save_npz(output_file, sparse_matrix)
        else:
            raise ValueError("Format must be 'numpy', 'csv', or 'sparse'")
        
        print(f"Token-document matrix saved to {output_file} in {format} format")
    
    def get_updated_corpus_data(self) -> List[Dict]:
        """Get the updated corpus data with context lengths included."""
        return self.chunks_data.copy()