import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import os
import struct
import heapq
import pickle
from collections import OrderedDict, defaultdict
import threading
import time
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import tkinter as tk

# ------------------------------
# 1. Huffman coding utilities
# ------------------------------
def build_huffman_tree(probs):
    """Build Huffman codes given symbol probabilities.
    Returns dict {symbol: code_str} and dict {symbol: code_int, length}
    """
    heap = [[wt, [sym, ""]] for sym, wt in probs.items() if wt > 0]
    heapq.heapify(heap)
    while len(heap) > 1:
        lo = heapq.heappop(heap)
        hi = heapq.heappop(heap)
        for pair in lo[1:]:
            pair[1] = '0' + pair[1]
        for pair in hi[1:]:
            pair[1] = '1' + pair[1]
        heapq.heappush(heap, [lo[0] + hi[0]] + lo[1:] + hi[1:])
    codes = {sym: code for sym, code in heap[0][1:]}
    return codes

def huffman_encode(symbols, codes):
    """Encode list of symbols into bytes using Huffman codes."""
    bitstring = ''.join(codes[s] for s in symbols)
    # pad to multiple of 8
    padding = 8 - (len(bitstring) % 8)
    if padding != 8:
        bitstring += '0' * padding
    else:
        padding = 0
    # convert to bytes
    byte_array = bytearray()
    for i in range(0, len(bitstring), 8):
        byte = int(bitstring[i:i+8], 2)
        byte_array.append(byte)
    return bytes(byte_array), padding

def huffman_decode(byte_data, codes_rev, num_symbols, padding):
    """Decode bytes back to symbols."""
    bitstring = ''.join(f'{b:08b}' for b in byte_data)
    if padding:
        bitstring = bitstring[:-padding]
    symbols = []
    cur = ''
    for bit in bitstring:
        cur += bit
        if cur in codes_rev:
            symbols.append(codes_rev[cur])
            cur = ''
    assert len(symbols) == num_symbols
    return symbols

# ------------------------------
# 2. Block splitting and compression
# ------------------------------
def split_into_blocks(tensor, block_size):
    """Split a 1D or 2D tensor into blocks of size <= block_size.
    Returns list of blocks (list of values).
    """
    flat = tensor.flatten()
    if hasattr(flat, 'cpu'):
        flat = flat.cpu().numpy()
    blocks = []
    for i in range(0, len(flat), block_size):
        blocks.append(flat[i:i+block_size].tolist())
    return blocks

def compress_block(values, block_id, output_dir):
    """Quantize block to ternary, Huffman encode, write to file.
    Returns file path and metadata (scale, shape, block_id).
    """
    # Ternary quantization: scale = mean abs of non‑zeros
    arr = np.array(values)
    nonzero = arr[np.abs(arr) > 1e-6]
    if len(nonzero) == 0:
        scale = 1.0
    else:
        scale = np.mean(np.abs(nonzero))
    thr = 0.5 * scale   # threshold for zero
    ternary = np.zeros_like(arr, dtype=np.int8)
    ternary[arr > thr] = 1
    ternary[arr < -thr] = -1
    # Symbol probabilities
    unique, counts = np.unique(ternary, return_counts=True)
    probs = {int(sym): cnt/len(ternary) for sym, cnt in zip(unique, counts)}
    codes = build_huffman_tree(probs)
    # Encode
    symbols = ternary.tolist()
    encoded_bytes, padding = huffman_encode(symbols, codes)
    # Save file
    file_path = os.path.join(output_dir, f"block_{block_id:06d}.bin")
    with open(file_path, 'wb') as f:
        # Header: scale (float32), num_symbols (int32), padding (uint8), codebook size (int32)
        f.write(struct.pack('f', scale))
        f.write(struct.pack('i', len(symbols)))
        f.write(struct.pack('B', padding))
        # Save codebook: {symbol: code_str}
        codebook_data = pickle.dumps(codes)
        f.write(struct.pack('i', len(codebook_data)))
        f.write(codebook_data)
        # Save encoded data
        f.write(encoded_bytes)
    return file_path, scale, block_id

def decompress_block(file_path):
    """Read a compressed block file and return the original values (floats)."""
    with open(file_path, 'rb') as f:
        scale = struct.unpack('f', f.read(4))[0]
        num_symbols = struct.unpack('i', f.read(4))[0]
        padding = struct.unpack('B', f.read(1))[0]
        codebook_size = struct.unpack('i', f.read(4))[0]
        codebook_data = f.read(codebook_size)
        codes = pickle.loads(codebook_data)
        encoded_bytes = f.read()
    # Build reverse mapping
    rev_codes = {v: k for k, v in codes.items()}
    symbols = huffman_decode(encoded_bytes, rev_codes, num_symbols, padding)
    # Convert back to float values
    values = np.array(symbols, dtype=np.float32) * scale
    return values

# ------------------------------
# 3. Model with on‑demand block loading and visualizer
# ------------------------------
class BlockCompressedLinear(nn.Module):
    """Linear layer that reconstructs weights from compressed blocks on-demand.
    Tracks access counts per block for visualization.
    """
    def __init__(self, in_features, out_features, block_files, original_shape, bias=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.block_files = block_files
        self.original_shape = original_shape
        self.register_buffer('_dummy', torch.tensor(0))
        # LRU cache: block_idx -> (values, last_access_time)
        self.cache = {}
        self.max_cache_size = 50
        # For visualizer: track access counts per block
        self.access_counts = defaultdict(int)
        if bias is not None:
            self.bias = nn.Parameter(bias)
        else:
            self.bias = None

    def _load_block(self, block_idx):
        if block_idx in self.cache:
            self.cache[block_idx] = (self.cache[block_idx][0], time.time())
            return self.cache[block_idx][0]
        file_path = self.block_files[block_idx]
        values = decompress_block(file_path)
        self.cache[block_idx] = (values, time.time())
        if len(self.cache) > self.max_cache_size:
            oldest = min(self.cache.keys(), key=lambda k: self.cache[k][1])
            del self.cache[oldest]
        return values

    def forward(self, x):
        all_values = []
        for idx in range(len(self.block_files)):
            vals = self._load_block(idx)
            all_values.extend(vals)
            self.access_counts[idx] += 1
        weight = torch.tensor(all_values, device=x.device, dtype=torch.float32).reshape(self.original_shape)
        out = F.linear(x, weight, self.bias)
        return out


class BlockCompressedConv2d(nn.Module):
    """Conv2d layer that reconstructs weights from compressed blocks on-demand.
    Tracks access counts per block for visualization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, block_files, original_shape, padding=0, bias=None):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        if isinstance(kernel_size, int):
            self.kernel_size = (kernel_size, kernel_size)
        else:
            self.kernel_size = kernel_size
        self.padding = padding
        self.block_files = block_files
        self.original_shape = original_shape
        self.register_buffer('_dummy', torch.tensor(0))
        # LRU cache
        self.cache = {}
        self.max_cache_size = 50
        # For visualizer
        self.access_counts = defaultdict(int)
        if bias is not None:
            self.bias = nn.Parameter(bias)
        else:
            self.bias = None

    def _load_block(self, block_idx):
        if block_idx in self.cache:
            self.cache[block_idx] = (self.cache[block_idx][0], time.time())
            return self.cache[block_idx][0]
        file_path = self.block_files[block_idx]
        values = decompress_block(file_path)
        self.cache[block_idx] = (values, time.time())
        if len(self.cache) > self.max_cache_size:
            oldest = min(self.cache.keys(), key=lambda k: self.cache[k][1])
            del self.cache[oldest]
        return values

    def forward(self, x):
        all_values = []
        for idx in range(len(self.block_files)):
            vals = self._load_block(idx)
            all_values.extend(vals)
            self.access_counts[idx] += 1
        weight = torch.tensor(all_values, device=x.device, dtype=torch.float32).reshape(self.original_shape)
        out = F.conv2d(x, weight, self.bias, padding=self.padding)
        return out

# ------------------------------
# 4. Build compressed model from trained model
# ------------------------------
def build_compressed_model(original_model, block_size, output_dir):
    """Take a trained model, compress all weights into block files,
    and return a new CompressedModel that performs real forward passes
    with on-demand block loading.
    """
    os.makedirs(output_dir, exist_ok=True)
    mapping, total_blocks = compress_entire_model(original_model, block_size, output_dir)

    # Build the compressed model
    comp_model = CompressedModel(original_model, mapping)
    return comp_model, mapping, total_blocks


class CompressedModel(nn.Module):
    """Wraps an original model's architecture but loads weights from compressed block files.
    Each layer (Conv2d or Linear) is replaced by its block-compressed counterpart.
    """
    def __init__(self, original_model, block_mapping):
        super().__init__()
        self.block_mapping = block_mapping  # layer_name -> list of block file paths
        # Build compressed layers
        self.conv1 = self._make_compressed_layer(original_model.conv1, 'conv1.weight')
        self.conv2 = self._make_compressed_layer(original_model.conv2, 'conv2.weight')
        self.fc1 = self._make_compressed_layer(original_model.fc1, 'fc1.weight')
        self.fc2 = self._make_compressed_layer(original_model.fc2, 'fc2.weight')
        self.pool = original_model.pool

    def _make_compressed_layer(self, orig_layer, weight_name):
        block_files = self.block_mapping[weight_name]
        orig_shape = orig_layer.weight.shape
        if isinstance(orig_layer, nn.Conv2d):
            return BlockCompressedConv2d(
                orig_layer.in_channels, orig_layer.out_channels,
                orig_layer.kernel_size, block_files, orig_shape,
                padding=orig_layer.padding[0] if isinstance(orig_layer.padding, tuple) else orig_layer.padding,
                bias=orig_layer.bias
            )
        elif isinstance(orig_layer, nn.Linear):
            return BlockCompressedLinear(
                orig_layer.in_features, orig_layer.out_features,
                block_files, orig_shape,
                bias=orig_layer.bias
            )
        else:
            raise TypeError(f"Unsupported layer type: {type(orig_layer)}")

    def get_all_access_counts(self):
        """Aggregate access counts from all compressed layers.
        Returns dict: global_block_id -> count
        """
        counts = {}
        global_id = 0
        for layer in [self.conv1, self.conv2, self.fc1, self.fc2]:
            for local_id in range(len(layer.block_files)):
                counts[global_id] = layer.access_counts.get(local_id, 0)
                global_id += 1
        return counts

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# ------------------------------
# 6. Visualizer (Tkinter + Matplotlib)
# ------------------------------
class BlockVisualizer:
    def __init__(self, num_blocks, compressed_model, test_loader):
        self.num_blocks = num_blocks
        self.compressed_model = compressed_model
        self.test_loader = test_loader
        self.root = tk.Tk()
        self.root.title("Block Activation Visualizer – Real Inference")
        self.fig, (self.ax_grid, self.ax_bar) = plt.subplots(1, 2, figsize=(16, 8))
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        self.running = True
        self.update_interval = 300  # ms
        self.images_processed = 0
        self._draw_grid()
        self.root.after(self.update_interval, self._update_from_inference)
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)

    def _draw_grid(self):
        self.ax_grid.clear()
        cols = int(np.ceil(np.sqrt(self.num_blocks)))
        rows = int(np.ceil(self.num_blocks / cols))
        for i in range(self.num_blocks):
            r = i // cols
            c = i % cols
            rect = patches.Rectangle((c, rows-1-r), 1, 1, facecolor='#1a1a2e', edgecolor='#333', linewidth=0.5)
            self.ax_grid.add_patch(rect)
        self.ax_grid.set_xlim(0, cols)
        self.ax_grid.set_ylim(0, rows)
        self.ax_grid.set_aspect('equal')
        self.ax_grid.axis('off')
        self.ax_grid.set_title("Block Activation Map\n(blue=cold, red=hot)")
        self.ax_bar.set_title("Access Distribution")
        self.canvas.draw()

    def _update_from_inference(self):
        if not self.running:
            return
        # Run a batch of real inference
        device = next(self.compressed_model.parameters()).device
        self.compressed_model.eval()
        batch_count = 0
        with torch.no_grad():
            for data, target in self.test_loader:
                data, target = data.to(device), target.to(device)
                _ = self.compressed_model(data)
                batch_count += 1
                self.images_processed += len(data)
                if batch_count >= 2:  # process 2 batches per update
                    break

        # Pull access counts from the model
        counts = self.compressed_model.get_all_access_counts()
        max_count = max(counts.values()) if counts else 1

        # Update grid heatmap
        cols = int(np.ceil(np.sqrt(self.num_blocks)))
        rows = int(np.ceil(self.num_blocks / cols))
        self.ax_grid.clear()
        for i in range(self.num_blocks):
            r = i // cols
            c = i % cols
            count = counts.get(i, 0)
            intensity = count / max_count if max_count > 0 else 0
            # Interpolate from blue (cold) to red (hot)
            color = (intensity, 0, 1 - intensity)
            rect = patches.Rectangle((c, rows-1-r), 1, 1, facecolor=color, edgecolor='#333', linewidth=0.5)
            self.ax_grid.add_patch(rect)
        self.ax_grid.set_xlim(0, cols)
        self.ax_grid.set_ylim(0, rows)
        self.ax_grid.set_aspect('equal')
        self.ax_grid.axis('off')
        self.ax_grid.set_title(f"Block Activation Map\n{self.images_processed} images processed")

        # Update bar chart (top-30 most accessed blocks)
        self.ax_bar.clear()
        sorted_blocks = sorted(counts.items(), key=lambda x: x[1], reverse=True)[:30]
        if sorted_blocks:
            block_ids, acc_counts = zip(*sorted_blocks)
            self.ax_bar.barh(range(len(block_ids)), acc_counts, color='steelblue')
            self.ax_bar.set_yticks(range(len(block_ids)))
            self.ax_bar.set_yticklabels([f"#{b}" for b in block_ids])
            self.ax_bar.invert_yaxis()
            self.ax_bar.set_xlabel("Access Count")
            self.ax_bar.set_title("Top-30 Most Accessed Blocks")

        self.canvas.draw()
        self.root.after(self.update_interval, self._update_from_inference)

    def on_close(self):
        self.running = False
        self.root.destroy()

# ------------------------------
# 7. Main: train, compress, visualize
# ------------------------------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        #self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        #self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        #self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 128)
        self.fc4 = nn.Linear(128, 128)
        self.fc5 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = self.fc5(x)
        return x

def train_simple_cnn():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_loader = DataLoader(datasets.MNIST('../data', train=True, download=True, transform=transform), batch_size=64, shuffle=True)
    model = SimpleCNN().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    for epoch in range(1, 4):
        model.train()
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            out = model(data)
            loss = F.cross_entropy(out, target)
            loss.backward()
            optimizer.step()
    return model

def compress_entire_model(model, block_size, output_dir):
    """Compress all weight tensors into block files and return mapping layer->list of file paths."""
    os.makedirs(output_dir, exist_ok=True)
    mapping = {}
    global_block_id = 0
    for name, param in model.named_parameters():
        if 'weight' not in name:
            continue
        flat = param.data.cpu().numpy().flatten()
        blocks = split_into_blocks(flat, block_size)
        block_files = []
        for block_vals in blocks:
            file_path, _, _ = compress_block(block_vals, global_block_id, output_dir)
            block_files.append(file_path)
            global_block_id += 1
        mapping[name] = block_files
    # Also save a manifest
    with open(os.path.join(output_dir, 'manifest.pkl'), 'wb') as f:
        pickle.dump(mapping, f)
    return mapping, global_block_id

def run_visualization(compressed_model, test_loader, num_blocks):
    vis = BlockVisualizer(num_blocks, compressed_model, test_loader)
    vis.root.mainloop()

if __name__ == "__main__":
    # Train a model
    print("Training model...")
    model = train_simple_cnn()
    model.eval()

    # Compress it into small blocks and build the compressed model
    # Calculate block size to get approximately 16 blocks total
    # SimpleCNN has: conv1(1*32*3*3=288) + conv2(32*64*3*3=18432) + fc1(64*7*7*128=401408) + fc2(128*10=1280) = 421408 weights
    # For 16 blocks: 421408 / 16 = 26338 weights per block
    block_size = 26338   # each block contains 26338 weights to get ~16 blocks
    output_dir = "./compressed_blocks"
    compressed_model, mapping, total_blocks = build_compressed_model(model, block_size, output_dir)
    print(f"Created {total_blocks} block files in {output_dir}")

    # Prepare test loader for real inference visualization
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    test_loader = DataLoader(
        datasets.MNIST('../data', train=False, download=True, transform=transform),
        batch_size=32, shuffle=True
    )

    # Launch visualizer with REAL inference accesses
    print(f"Launching visualizer with real MNIST inference ({total_blocks} blocks)...")
    run_visualization(compressed_model, test_loader, total_blocks)