import numpy as np
import os
import math
from pathlib import Path
import matplotlib.pyplot as plt
import seaborn as sns

# Path settings
input_file_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\1Mb.txt"
output_directory = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\True Random Numbers Verification\02_heatmap_and_shannon_entropy_analysis"
output_result_filename = "bitstream_Shannon_entropy_analysis.txt"
output_image_filename = "bitstream_Shannon_entropy_heatmap.tif"
output_file_path = os.path.join(output_directory, output_result_filename)
output_image_path = os.path.join(output_directory, output_image_filename)

def calculate_bitstream_entropy(bitstream_file, n=10000000):
    if not Path(bitstream_file).exists():
        print(f"Error: File {bitstream_file} does not exist")
        return None
    
    try:
        with open(bitstream_file, 'r') as f:
            content = f.read().strip()
            bits = ''.join(c for c in content if c in '01')
            
            if len(bits) < n:
                print(f"Warning: File {bitstream_file} has fewer than {n} bits")
                bits = bits[:len(bits)]
            else:
                bits = bits[:n]
            
            count_1 = bits.count('1')
            count_0 = len(bits) - count_1
            
            p_1 = count_1 / len(bits) if len(bits) > 0 else 0
            p_0 = count_0 / len(bits) if len(bits) > 0 else 0
            
            entropy = 0
            if p_1 > 0:
                entropy -= p_1 * math.log2(p_1)
            if p_0 > 0:
                entropy -= p_0 * math.log2(p_0)
            
            return {
                'file_name': os.path.basename(bitstream_file),
                'entropy': entropy,
                'p_1': p_1,
                'p_0': p_0,
                'count_1': count_1,
                'count_0': count_0,
                'total_bits': len(bits)
            }
    except Exception as e:
        print(f"Error processing file {bitstream_file}: {e}")
        return None

def process_all_bitstreams(directory_path, output_file_path, n=1000000):
    directory = Path(directory_path)
    if not directory.exists():
        print(f"Error: Directory {directory_path} does not exist")
        return []
    
    bitstream_files = list(directory.glob("*.txt"))
    if not bitstream_files:
        print(f"Warning: No txt files found in directory {directory_path}")
        return []
    
    results = []
    print(f"Processing {len(bitstream_files)} bitstream files...")
    
    for i, file_path in enumerate(bitstream_files, 1):
        if i % 100 == 0 or i == len(bitstream_files):
            print(f"Progress: {i}/{len(bitstream_files)}")
        
        result = calculate_bitstream_entropy(file_path, n)
        if result:
            results.append(result)
    
    save_entropy_results(results, output_file_path, n)
    
    return results

def save_entropy_results(results, output_path, n):
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("Bitstream Shannon Entropy Analysis Results (with Detailed Calculation Process)\n")
        f.write("="*60 + "\n\n")
        
        f.write("[Calculation Principle]\n")
        f.write("Information entropy formula: H = -Σ p(i) * log₂(p(i))\n")
        f.write("Where p(i) is the probability of value i in the bitstream, i is 0 or 1\n")
        f.write("p(1) = Number of 1s / Total bitstream length\n")
        f.write("p(0) = Number of 0s / Total bitstream length\n\n")
        
        if results:
            avg_entropy = np.mean([r['entropy'] for r in results])
            avg_p1 = np.mean([r['p_1'] for r in results])
            max_entropy = max([r['entropy'] for r in results])
            min_entropy = min([r['entropy'] for r in results])
            
            f.write("[Overall Statistics]\n")
            f.write(f"Total files processed: {len(results)}\n")
            f.write(f"Bits per file: {n}\n")
            f.write(f"Average Shannon entropy: {avg_entropy:.6f} bits\n")
            f.write(f"Average probability of 1: {avg_p1:.6f}\n")
            f.write(f"Maximum Shannon entropy: {max_entropy:.6f} bits\n")
            f.write(f"Minimum Shannon entropy: {min_entropy:.6f} bits\n")
            f.write(f"Theoretical maximum entropy: 1.0 bits (when p(0)=p(1)=0.5)\n\n")
        
        if results:
            f.write("[Detailed Calculation Process Example]\n")
            f.write("="*60 + "\n\n")
            
            sample_results = results[:3]
            for idx, result in enumerate(sample_results, 1):
                f.write(f"Example {idx}: File {result['file_name']}\n")
                f.write("="*60 + "\n")
                f.write(f"File path: {result['file_name']}\n")
                f.write(f"Processed bits: {result['total_bits']}\n")
                f.write(f"Number of 1s: {result['count_1']}\n")
                f.write(f"Number of 0s: {result['count_0']}\n")
                f.write(f"Probability of 1 p(1) = {result['count_1']}/{result['total_bits']} = {result['p_1']:.6f}\n")
                f.write(f"Probability of 0 p(0) = {result['count_0']}/{result['total_bits']} = {result['p_0']:.6f}\n")
                f.write(f"Shannon entropy H = -[p(1)*log2(p(1)) + p(0)*log2(p(0))]\n")
                
                if result['p_1'] > 0:
                    term1 = -result['p_1'] * math.log2(result['p_1'])
                    f.write(f"  = -[{result['p_1']:.6f}*log2({result['p_1']:.6f}) + ")
                else:
                    term1 = 0
                    f.write(f"  = -[0 + ")
                
                if result['p_0'] > 0:
                    term2 = -result['p_0'] * math.log2(result['p_0'])
                    f.write(f"{result['p_0']:.6f}*log2({result['p_0']:.6f})]\n")
                else:
                    term2 = 0
                    f.write(f"0]\n")
                
                f.write(f"  = -[{term1:.6f} + {term2:.6f}]\n")
                f.write(f"  = {result['entropy']:.6f} bits\n\n")
        
        if results:
            f.write("[All File Results List]\n")
            f.write("="*60 + "\n")
            f.write(f"{'File Name':<20} {'Shannon Entropy':<15} {'Probability of 1':<15}\n")
            f.write("-"*60 + "\n")
            
            for result in results:
                f.write(f"{result['file_name']:<20} {result['entropy']:<15.6f} {result['p_1']:<15.6f}\n")
    
    print(f"Entropy analysis results saved to: {output_path}")

def save_single_entropy_result(result, output_path, n):
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("Bitstream Shannon Entropy Analysis Results (with Detailed Calculation Process)\n")
        f.write("="*60 + "\n\n")
        
        f.write("[Calculation Principle]\n")
        f.write("Information entropy formula: H = -Σ p(i) * log₂(p(i))\n")
        f.write("Where p(i) is the probability of value i in the bitstream, i is 0 or 1\n")
        f.write("p(1) = Number of 1s / Total bitstream length\n")
        f.write("p(0) = Number of 0s / Total bitstream length\n\n")
        
        f.write("[File Analysis Results]\n")
        f.write("="*60 + "\n\n")
        
        f.write(f"File name: {result['file_name']}\n")
        f.write(f"File path: {result['file_name']}\n")
        f.write(f"Processed bits: {result['total_bits']}\n")
        f.write(f"Number of 1s: {result['count_1']}\n")
        f.write(f"Number of 0s: {result['count_0']}\n")
        f.write(f"Probability of 1 p(1) = {result['count_1']}/{result['total_bits']} = {result['p_1']:.6f}\n")
        f.write(f"Probability of 0 p(0) = {result['count_0']}/{result['total_bits']} = {result['p_0']:.6f}\n")
        f.write(f"Shannon entropy H = -[p(1)*log2(p(1)) + p(0)*log2(p(0))]\n")
        
        if result['p_1'] > 0:
            term1 = -result['p_1'] * math.log2(result['p_1'])
            f.write(f"  = -[{result['p_1']:.6f}*log2({result['p_1']:.6f}) + ")
        else:
            term1 = 0
            f.write(f"  = -[0 + ")
        
        if result['p_0'] > 0:
            term2 = -result['p_0'] * math.log2(result['p_0'])
            f.write(f"{result['p_0']:.6f}*log2({result['p_0']:.6f})]\n")
        else:
            term2 = 0
            f.write(f"0]\n")
        
        f.write(f"  = -[{term1:.6f} + {term2:.6f}]\n")
        f.write(f"  = {result['entropy']:.6f} bits\n\n")
        
        f.write("[Overall Statistics]\n")
        f.write(f"Total files processed: 1\n")
        f.write(f"Bits per file: {n}\n")
        f.write(f"Shannon entropy: {result['entropy']:.6f} bits\n")
        f.write(f"Probability of 1: {result['p_1']:.6f}\n")
        f.write(f"Theoretical maximum entropy: 1.0 bits (when p(0)=p(1)=0.5)\n")
    
    print(f"Entropy analysis results saved to: {output_path}")

def generate_bitstream_visualization(bitstream_file, output_image_path, n=10000, rows=100, cols=100):
    if not Path(bitstream_file).exists():
        print(f"Error: File {bitstream_file} does not exist")
        return None
    
    try:
        with open(bitstream_file, 'r') as f:
            content = f.read().strip()
            bits = ''.join(c for c in content if c in '01')
            
            if len(bits) < n:
                print(f"Warning: File {bitstream_file} has fewer than {n} bits")
                bits = bits[:len(bits)]
            else:
                bits = bits[:n]
            
            count_1 = bits.count('1')
            count_0 = len(bits) - count_1
            p_1 = count_1 / len(bits) if len(bits) > 0 else 0
            
            entropy = 0
            if p_1 > 0:
                entropy -= p_1 * math.log2(p_1)
            p_0 = 1 - p_1
            if p_0 > 0:
                entropy -= p_0 * math.log2(p_0)
            
            grid_data = np.zeros((rows, cols))
            
            for i in range(min(n, len(bits))):
                row = i // cols
                col = i % cols
                grid_data[row, col] = int(bits[i])
            
            output_dir = os.path.dirname(output_image_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            plt.figure(figsize=(10, 8), dpi=300)
            ax = plt.gca()
            ax.imshow(grid_data, cmap='binary', aspect='equal', interpolation='nearest')
            ax.invert_yaxis()
            
            plt.title(f'Shannon entropy = {entropy:.6f}  P={p_1*100:.2f}%', 
                      fontsize=30, fontname='Arial', pad=30)
            
            xticks = np.arange(0, cols + 1, 200)
            yticks = np.arange(rows, -1, -200)
            plt.xticks(xticks, fontsize=26)
            plt.yticks(yticks, fontsize=26)
            
            plt.grid(alpha=0.3, color='gray', linestyle='-', linewidth=0.5)
            
            plt.tight_layout()
            plt.savefig(output_image_path, dpi=600, bbox_inches='tight')
            plt.close()
            
            print(f"Bitstream visualization heatmap saved to: {output_image_path}")
            return {
                'file_name': os.path.basename(bitstream_file),
                'entropy': entropy,
                'p_1': p_1,
                'count_1': count_1,
                'count_0': count_0,
                'total_bits': len(bits)
            }
    except Exception as e:
        print(f"Error generating bitstream visualization: {e}")
        return None

def main():
    print("===== Bitstream Shannon Entropy Analysis Configuration =====")
    
    print(f"Input bitstream file: {input_file_path}")
    
    if not os.path.exists(input_file_path):
        print(f"Error: File {input_file_path} does not exist!")
        return
    
    with open(input_file_path, 'r') as f:
        content = f.read().strip()
        bits = ''.join(c for c in content if c in '01')
        total_bits_available = len(bits)
    
    print(f"Output result file: {output_file_path}")
    print(f"Output visualization heatmap: {output_image_path}")
    
    side_length = int(math.sqrt(total_bits_available))
    rows = side_length
    cols = side_length
    n = rows * cols
    
    print(f"Total bitstream length: {total_bits_available}")
    print(f"Heatmap configuration: {rows} rows × {cols} columns = {n} bits")
    print("="*50)
    
    print(f"Processing file: {os.path.basename(input_file_path)}")
    result = calculate_bitstream_entropy(input_file_path, n)
    
    if result:
        save_single_entropy_result(result, output_file_path, n)
        
        visualization_result = generate_bitstream_visualization(input_file_path, output_image_path, n, rows, cols)
        
        print("\nBitstream Shannon entropy analysis completed!")
        print(f"Shannon entropy: {result['entropy']:.6f} bits")
        print(f"Probability of 1: {result['p_1']*100:.2f}%")
        print(f"Processed bits: {result['total_bits']}")
        print(f"Number of 1s: {result['count_1']}")
        print(f"Number of 0s: {result['count_0']}")
    else:
        print("File processing failed")

if __name__ == "__main__":
    main()