import sqlite3
import struct
import mpmath
from pathlib import Path
import glob
import re
import numpy as np
import time

# Set precision for mpmath
mpmath.mp.prec = 300
eps = mpmath.mpf(2) ** (-101)

# ========== PLATT ZERO READER ==========

class PlattZeroReader:
    def __init__(self, data_folder="sft_output"):
        self.data_folder = Path(data_folder)
        self.db_path = self.data_folder / "index.db"
        self.conn = None
        
    def connect(self):
        if not self.db_path.exists():
            raise FileNotFoundError(f"index.db not found in {self.data_folder}")
        self.conn = sqlite3.connect(str(self.db_path))
        
    def get_starting_point(self, N):
        cursor = self.conn.cursor()
        cursor.execute("""
            SELECT filename, offset, block_number, t, N
            FROM zero_index
            WHERE N <= ?
            ORDER BY N DESC
            LIMIT 1
        """, (N,))
        result = cursor.fetchone()
        if result is None:
            raise ValueError(f"No data found for N={N}")
        return result
    
    def zeros_starting_at_N(self, start_n, count):
        self.connect()
        
        filename, offset, block_number, t0, N0 = self.get_starting_point(start_n)
        
        filepath = self.data_folder / filename
        if not filepath.exists():
            raise FileNotFoundError(f"Binary file not found: {filepath}")
        
        infile = open(filepath, 'rb')
        num_blocks = struct.unpack('Q', infile.read(8))[0]
        infile.seek(offset)
        
        t0, t1, Nt0, Nt1 = struct.unpack('ddQQ', infile.read(8 * 4))
        
        mpmath.mp.prec = max(300, int(np.log2(t1)) + 10 + 101)
        t0_mpf = mpmath.mpf(t0)
        
        Z = 0
        n = Nt0
        
        while n < start_n and n < Nt1:
            z1, z2, z3 = struct.unpack('QIB', infile.read(13))
            Z = Z + (z3 << 96) + (z2 << 64) + z1
            n += 1
        
        produced = 0
        while produced < count:
            if n == Nt1:
                block_number += 1
                if block_number == num_blocks:
                    infile.close()
                    cursor = self.conn.cursor()
                    cursor.execute("""
                        SELECT filename, offset, block_number, t, N
                        FROM zero_index
                        WHERE N >= ?
                        ORDER BY N ASC
                        LIMIT 1
                    """, (n,))
                    result = cursor.fetchone()
                    if result is None:
                        break
                    filename, offset, block_number, t0, N0 = result
                    filepath = self.data_folder / filename
                    infile = open(filepath, 'rb')
                    num_blocks = struct.unpack('Q', infile.read(8))[0]
                    infile.seek(offset)
                
                header = infile.read(8 * 4)
                t0, t1, Nt0, Nt1 = struct.unpack('ddQQ', header)
                mpmath.mp.prec = max(300, int(np.log2(t1)) + 10 + 101)
                t0_mpf = mpmath.mpf(t0)
                Z = 0
            
            z1, z2, z3 = struct.unpack('QIB', infile.read(13))
            Z = Z + (z3 << 96) + (z2 << 64) + z1
            
            gamma = float(t0_mpf + mpmath.mpf(Z) * eps)
            
            yield n, gamma
            
            n += 1
            produced += 1
        
        infile.close()
        self.conn.close()


# ========== SFT CHUNK READER ==========

def load_sft_data(chunk_pattern, target_start, target_end):
    files = glob.glob(chunk_pattern)
    print(f"Found {len(files)} chunk files")
    
    sft_dict = {}
    for filepath in sorted(files):
        with open(filepath, 'r') as f:
            for line in f:
                if line.startswith('#') or not line.strip():
                    continue
                parts = line.split()
                if len(parts) >= 3:
                    n = int(parts[0])
                    if target_start <= n <= target_end:
                        sft_dict[n] = float(parts[2])
    
    print(f"  Loaded {len(sft_dict)} SFT entries in range")
    return sft_dict


# ========== COMPARISON WITH RELATIVE ERROR ==========

def compare_streaming(sft_dict, zeros_iter, start_n, end_n):
    print("\nComparing zeros (absolute and relative error)...")
    
    count = 0
    sum_abs_err = 0.0
    sum_rel_err = 0.0
    sum_abs_err2 = 0.0
    sum_rel_err2 = 0.0
    max_abs_err = 0.0
    max_rel_err = 0.0
    min_rel_err = float('inf')
    errors = []
    rel_errors = []
    
    last_report = 0
    start_time = time.time()
    
    for n, gamma in zeros_iter:
        if n in sft_dict:
            e_sft = sft_dict[n]
            abs_err = e_sft - gamma
            rel_err = abs(abs_err / gamma) if gamma != 0 else 0
            
            count += 1
            sum_abs_err += abs(abs_err)
            sum_rel_err += rel_err
            sum_abs_err2 += abs_err * abs_err
            sum_rel_err2 += rel_err * rel_err
            
            if abs(abs_err) > max_abs_err:
                max_abs_err = abs(abs_err)
            if rel_err > max_rel_err:
                max_rel_err = rel_err
            if rel_err < min_rel_err:
                min_rel_err = rel_err
            
            errors.append(abs_err)
            rel_errors.append(rel_err)
            if len(errors) > 1000:
                errors.pop(0)
                rel_errors.pop(0)
            
            if count % 10000 == 0 and count != last_report:
                last_report = count
                elapsed = time.time() - start_time
                rate = count / elapsed if elapsed > 0 else 0
                mean_abs_err = sum_abs_err / count
                mean_rel_err = sum_rel_err / count
                print(f"\n--- After {count:,} matches (n={n:,}, {rate:.0f}/s) ---")
                print(f"  Mean absolute error: {mean_abs_err:+.6f}")
                print(f"  Mean relative error: {mean_rel_err:.6e}")
                print(f"  Max absolute error:  {max_abs_err:.6f}")
                print(f"  Max relative error:  {max_rel_err:.6e}")
    
    return {
        'count': count,
        'sum_abs_err': sum_abs_err,
        'sum_rel_err': sum_rel_err,
        'sum_abs_err2': sum_abs_err2,
        'sum_rel_err2': sum_rel_err2,
        'max_abs_err': max_abs_err,
        'max_rel_err': max_rel_err,
        'min_rel_err': min_rel_err,
        'errors': errors,
        'rel_errors': rel_errors
    }


# ========== MAIN ==========

def main():
    print("\n" + "="*60)
    print("SFT vs RIEMANN ZERO COMPARISON (with Relative Error)")
    print("="*60)
    
    data_folder = "sft_output"
    chunk_pattern = "sft_output/sft_kreal_n*_chunk*.txt"
    
    start_n = int(input("Enter starting n: "))
    end_n = int(input("Enter ending n: "))
    
    print("\nLoading SFT data...")
    sft_dict = load_sft_data(chunk_pattern, start_n, end_n)
    
    if not sft_dict:
        print("No SFT data found for range!")
        return
    
    print("\nReading zeros from Platt database...")
    reader = PlattZeroReader(data_folder)
    
    try:
        zeros_iter = reader.zeros_starting_at_N(start_n, end_n - start_n + 1)
    except Exception as e:
        print(f"Error reading zeros: {e}")
        print("\nMake sure these files are in sft_output/:")
        print("  - index.db")
        print("  - zeros_*.dat (all the binary files)")
        return
    
    results = compare_streaming(sft_dict, zeros_iter, start_n, end_n)
    
    count = results['count']
    
    if count == 0:
        print("No matching zeros found!")
        return
    
    mean_abs_err = results['sum_abs_err'] / count
    mean_rel_err = results['sum_rel_err'] / count
    rms_abs_err = np.sqrt(results['sum_abs_err2'] / count)
    rms_rel_err = np.sqrt(results['sum_rel_err2'] / count)
    std_abs_err = np.sqrt(results['sum_abs_err2'] / count - mean_abs_err**2)
    std_rel_err = np.sqrt(results['sum_rel_err2'] / count - mean_rel_err**2)
    
    median_abs_err = np.median(results['errors'])
    median_rel_err = np.median(results['rel_errors'])
    
    print("\n" + "="*60)
    print("FINAL RESULTS")
    print("="*60)
    print(f"Range: n={start_n:,} to {end_n:,}")
    print(f"Matched points: {count:,}")
    print("-"*40)
    print("ABSOLUTE ERROR (E_SFT - gamma_n):")
    print(f"  Mean error:     {mean_abs_err:+.6f}")
    print(f"  Median error:   {median_abs_err:+.6f}")
    print(f"  Std error:      {std_abs_err:.6f}")
    print(f"  RMS error:      {rms_abs_err:.6f}")
    print(f"  Max |error|:    {results['max_abs_err']:.6f}")
    print("-"*40)
    print("RELATIVE ERROR (|E_SFT - gamma_n| / gamma_n):")
    print(f"  Mean rel error: {mean_rel_err:.6e}")
    print(f"  Median rel err: {median_rel_err:.6e}")
    print(f"  Std rel error:  {std_rel_err:.6e}")
    print(f"  RMS rel error:  {rms_rel_err:.6e}")
    print(f"  Max rel error:  {results['max_rel_err']:.6e}")
    print(f"  Min rel error:  {results['min_rel_err']:.6e}")
    print("="*60)
    
    save = input("\nSave results to file? (y/n): ").lower().strip()
    if save == 'y':
        outfile = f"comparison_rel_{start_n}_{end_n}.txt"
        with open(outfile, 'w') as f:
            f.write(f"SFT vs Riemann Zero Comparison (with Relative Error)\n")
            f.write(f"Range: n={start_n} to {end_n}\n")
            f.write(f"Points: {count}\n\n")
            f.write(f"Absolute Error (E_SFT - gamma_n):\n")
            f.write(f"  Mean: {mean_abs_err:+.6f}\n")
            f.write(f"  Median: {median_abs_err:+.6f}\n")
            f.write(f"  Std: {std_abs_err:.6f}\n")
            f.write(f"  RMS: {rms_abs_err:.6f}\n")
            f.write(f"  Max |error|: {results['max_abs_err']:.6f}\n\n")
            f.write(f"Relative Error (|E_SFT - gamma_n| / gamma_n):\n")
            f.write(f"  Mean: {mean_rel_err:.6e}\n")
            f.write(f"  Median: {median_rel_err:.6e}\n")
            f.write(f"  Std: {std_rel_err:.6e}\n")
            f.write(f"  RMS: {rms_rel_err:.6e}\n")
            f.write(f"  Max: {results['max_rel_err']:.6e}\n")
            f.write(f"  Min: {results['min_rel_err']:.6e}\n")
        print(f"Saved to {outfile}")

if __name__ == "__main__":
    main()