quantify_all_metabolites_v3_ref20.py

#!/usr/bin/env python3
"""
Absolute Quantification of All Metabolites using Physics-Based Multi-Region Fitting

This version uses FILE 20 as reference instead of FILE 10, and excludes FILE 10 entirely.

Key improvements in v3:
1. Multi-region fitting: Fits ALL distinct chemical shift regions for each metabolite
2. Proton-weighted averaging: Combines region results weighted by proton count
3. Physics-based peak counting: Number of peaks = number of distinct proton environments
4. Internal consistency validation: Compare results from different regions

MODIFIED: Uses file 20 as reference, excludes file 10
"""

import nmrglue as ng
import numpy as np
from scipy.ndimage import gaussian_filter1d
from scipy.signal import find_peaks
from lmfit import Model
from lmfit.models import LorentzianModel, ConstantModel
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd
import os


def read_and_process(filepath):
    """Read JCAMP-DX file and return ppm, magnitude spectrum"""
    dic, data = ng.jcampdx.read(filepath)
    data_real = data[0]
    data_imag = data[1] if len(data) > 1 else [0] * len(data[0])
    magnitude = np.sqrt(data_real**2 + np.array(data_imag)**2)
    smooth = gaussian_filter1d(magnitude, sigma=2)
    
    sfo1 = float(dic['$SFO1'][0])
    o1_hz = float(dic['$O1'][0])
    sw_hz = float(dic['$SWH'][0])
    o1_ppm = o1_hz / sfo1
    sw_ppm = sw_hz / sfo1
    ppm = np.linspace(o1_ppm + sw_ppm/2, o1_ppm - sw_ppm/2, len(magnitude))
    return ppm, smooth


def find_tsp_peak(ppm, intensity):
    """Find TSP reference peak position"""
    mask = (ppm >= -0.5) & (ppm <= 0.5)
    return ppm[mask][np.argmax(intensity[mask])]


def integrate_peak(ppm, intensity, region):
    """Integrate a spectral region"""
    mask = (ppm >= region[0]) & (ppm <= region[1])
    return abs(np.trapz(intensity[mask], ppm[mask]))


def detect_peaks_in_region(ppm, intensity, region, expected_centers, n_peaks=1):
    """
    Detect peaks in spectral region and match to expected positions.
    Returns detected peak positions closest to expected centers.
    """
    mask = (ppm >= region[0]) & (ppm <= region[1])
    x_region = ppm[mask]
    y_region = intensity[mask]
    
    if len(x_region) == 0 or np.max(y_region) <= 0:
        return expected_centers[:n_peaks]
    
    # Find peaks with adaptive threshold
    max_intensity = np.max(y_region)
    # Use lower threshold (5%) to catch peaks at low concentrations and pH-shifted peaks
    height_threshold = max_intensity * 0.05
    min_distance = len(x_region) // (n_peaks * 3)
    
    peaks_idx, _ = find_peaks(y_region, height=height_threshold, distance=max(10, min_distance))
    
    if len(peaks_idx) == 0:
        return expected_centers[:n_peaks]
    
    peak_positions = x_region[peaks_idx]
    peak_heights = y_region[peaks_idx]
    
    # Match expected centers to detected peaks using one-to-one assignment
    # Priority: tallest peak first (intensity over proximity)
    tolerance = 0.30  # 0.3 ppm tolerance for pH-shifting metabolites
    
    detected = []
    used_peak_idx = set()
    
    for expected in expected_centers[:n_peaks]:
        # Find all peaks within tolerance that haven't been used
        distances = np.abs(peak_positions - expected)
        within_tolerance = (distances < tolerance) & ~np.isin(np.arange(len(peak_positions)), list(used_peak_idx))
        
        if np.any(within_tolerance):
            # Among available peaks within tolerance, pick the TALLEST
            valid_indices = np.where(within_tolerance)[0]
            valid_heights = peak_heights[valid_indices]
            tallest_local_idx = np.argmax(valid_heights)
            best_idx = valid_indices[tallest_local_idx]
            detected.append(peak_positions[best_idx])
            used_peak_idx.add(best_idx)
        else:
            # No unused peak found within tolerance, use expected position
            detected.append(expected)
    
    return detected


def get_metabolite_info_v3():
    """
    Return metabolite information with ALL distinct proton environments.
    
    Each metabolite defines:
    - regions: List of (start_ppm, end_ppm) tuples for each distinct region
    - region_peaks: List of expected peak centers for each region
    - region_protons: Number of protons contributing to each region
    - region_names: Names of proton environments in each region
    """
    
    metabolites = {
        'Alanine': {
            'folder': 'Alanine-Reference',
            # Alanine: CH3-CH(NH2)-COOH
            # CH3 at 1.48 ppm (3H), CH at 3.78 ppm (1H)
            'regions': [
                (1.40, 1.55),   # CH3 region
                (3.70, 3.85)    # α-CH region  
            ],
            'region_peaks': [
                [1.48],         # CH3
                [3.79]          # α-CH
            ],
            'region_protons': [3, 1],
            'region_names': ['CH3', 'α-CH'],
            'ref_conc': 40.633068,
            'files': {10: 40.633068, 20: 20.316534, 30: 10.158267, 40: 5.079133, 
                     50: 2.539567, 60: 1.269783, 70: 0.634892}
        },
        
        'Valine': {
            'folder': 'Valine-Reference',
            # Valine: (CH3)2-CH-CH(NH2)-COOH
            # Two CH3 at 0.99 and 1.04 ppm (6H total), CH at 3.62 ppm (1H)
            'regions': [
                (0.90, 1.15),   # Two CH3 groups
                (3.55, 3.70)    # α-CH
            ],
            'region_peaks': [
                [0.99, 1.04],   # Two CH3 doublets
                [3.62]          # α-CH
            ],
            'region_protons': [6, 1],
            'region_names': ['(CH3)2', 'α-CH'],
            'ref_conc': 5.021349,
            'files': {10: 5.021349, 20: 2.251067, 30: 1.255337, 40: 0.627669,
                     50: 0.313834, 60: 0.156917, 70: 0.078459, 80: 0.039229}
        },
        
        'Arginine': {
            'folder': 'Arginine-Reference',
            # Arginine: H2N-C(=NH)-NH-(CH2)3-CH(NH2)-COOH
            # δ-CH2 at ~3.25 ppm (2H), γ-CH2 at ~1.67 ppm (2H)
            # β-CH2 at ~1.90 ppm (2H), α-CH at ~3.75 ppm (1H)
            'regions': [
                (1.58, 2.00),   # β-CH2 + γ-CH2 (4H total, 2 peaks)
                (3.15, 3.45),   # δ-CH2 (2H)
                (3.60, 3.90)    # α-CH (1H)
            ],
            'region_peaks': [
                [1.67, 1.90],   # γ-CH2, β-CH2
                [3.25],         # δ-CH2
                [3.75]          # α-CH
            ],
            'region_protons': [4, 2, 1],
            'region_names': ['β+γ-CH2', 'δ-CH2', 'α-CH'],
            'ref_conc': 16.696725,
            'files': {10: 16.696725, 20: 8.348363, 30: 4.174181, 40: 2.087091,
                     50: 1.043545, 60: 0.521773, 70: 0.260886}
        },
        
        'Lactate': {
            'folder': 'Lactate-Reference',
            # Lactate: CH3-CH(OH)-COOH
            # CH3 at 1.33 ppm (3H), CH at 4.12 ppm (1H)
            'regions': [
                (1.20, 1.45),   # CH3
                (4.05, 4.20)    # CH-OH
            ],
            'region_peaks': [
                [1.33],
                [4.12]
            ],
            'region_protons': [3, 1],
            'region_names': ['CH3', 'CH-OH'],
            'ref_conc': 97.685764,
            'files': {10: 97.685764, 20: 1.563401, 30: 3.052680, 40: 6.105360,
                     50: 12.210720, 60: 24.421441, 70: 48.842882}
        },
        
        'Glucose': {
            'folder': 'Glucose-Reference',
            # Glucose: Multiple CH and CH2 groups
            # The anomeric region (5.15-5.35 ppm) has a clean single peak at 5.24 ppm
            # The H-2 region (3.4-3.6 ppm) is too crowded with overlapping peaks
            # Use only the anomeric region for reliable quantification
            'regions': [
                (5.15, 5.35),   # Anomeric H-1 (clean single peak)
            ],
            'region_peaks': [
                [5.24],
            ],
            'region_protons': [1],
            'region_names': ['H-1 (anomeric)'],
            'ref_conc': 103.144654,
            'files': {10: 103.144654, 20: 51.572327, 30: 25.786164, 40: 12.893082,
                     50: 6.446541, 60: 3.223270, 70: 1.611635}
        },
        
        'Glutamate': {
            'folder': 'Glutamate-Reference',
            # Glutamate: HOOC-CH2-CH2-CH(NH2)-COOH
            # β-CH2 has two visible peaks (~2.07, ~2.14) forming a multiplet
            # γ-CH2 at ~2.42 ppm
            'regions': [
                (2.00, 2.50),   # β-CH2 + γ-CH2 region
            ],
            'region_peaks': [
                [2.07, 2.14, 2.42],   # 3 peaks: β-CH2 doublet + γ-CH2
            ],
            'region_protons': [4],
            'region_names': ['β-CH2 (2) + γ-CH2'],
            'ref_conc': 10.958532,
            'files': {10: 10.958532, 20: 5.479266, 30: 2.739633, 40: 1.369816,
                     50: 0.684908, 60: 0.342454, 70: 0.171227}
        },
        
        'Aspartate': {
            'folder': 'Aspartate-Reference',
            # Aspartate: HOOC-CH2-CH(NH2)-COOH
            # β-CH2 AB system at ~2.72 and ~2.85 ppm (2H total, 2 peaks)
            'regions': [
                (2.55, 2.95),   # β-CH2 AB system
            ],
            'region_peaks': [
                [2.715, 2.850],
            ],
            'region_protons': [2],
            'region_names': ['β-CH2 (AB)'],
            'ref_conc': 5.052592,
            'files': {10: 5.052592, 20: 2.526296, 30: 1.263148, 40: 0.631574,
                     50: 0.315787, 60: 0.157894, 70: 0.078947, 80: 0.039473}
        },
        
        'Isoleucine': {
            'folder': 'Isoleucine-Reference',
            # Isoleucine: CH3-CH2-CH(CH3)-CH(NH2)-COOH
            # Two CH3 groups at 0.94 and 1.00 ppm (6H total), α-CH at 3.68 ppm
            'regions': [
                (0.90, 1.05),   # Two CH3 groups
                (3.60, 3.80),   # α-CH
            ],
            'region_peaks': [
                [0.94, 1.00],
                [3.68],
            ],
            'region_protons': [6, 1],
            'region_names': ['(CH3)2', 'α-CH'],
            'ref_conc': 8.0,
            'files': {10: 8.0, 20: 3.871951, 30: 1.935976, 40: 0.967988,
                     50: 0.483994, 60: 0.241997, 70: 0.120999}
        },
        
        'Leucine': {
            'folder': 'Leucine-Reference',
            # Leucine: (CH3)2-CH-CH2-CH(NH2)-COOH
            # Two diastereotopic CH3 groups at ~0.96 and ~0.98 ppm (6H total, 2 peaks)
            # The two methyls on the isopropyl group are diastereotopic due to the chiral center
            # α-CH at ~3.74 ppm (1H)
            'regions': [
                (0.90, 1.05),   # Two diastereotopic CH3 groups
                (3.65, 3.85),   # α-CH
            ],
            'region_peaks': [
                [0.959, 0.975],  # Two methyl doublets (diastereotopic)
                [3.74],          # α-CH
            ],
            'region_protons': [6, 1],
            'region_names': ['(CH3)2', 'α-CH'],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.743902, 30: 1.371951, 40: 0.685976,
                     50: 0.342988, 60: 0.171494, 70: 0.085747}
        },
        
        'Glutamine': {
            'folder': 'Glutamine-Reference',
            # Glutamine: H2N-C(=O)-CH2-CH2-CH(NH2)-COOH
            # γ-CH2 at ~2.14 ppm (2H), β-CH2 at ~2.46 ppm (2H), α-CH at ~3.78 ppm (1H)
            'regions': [
                (2.05, 2.60),   # β-CH2 + γ-CH2 (side chain methylenes)
                (3.70, 3.90),   # α-CH
            ],
            'region_peaks': [
                [2.14, 2.46],
                [3.78],
            ],
            'region_protons': [4, 1],
            'region_names': ['β+γ-CH2', 'α-CH'],
            'ref_conc': 17.0,
            'files': {10: 17.0, 20: 8.477011, 30: 4.238505, 40: 2.119253,
                     50: 1.059626, 60: 0.529813, 70: 0.264906}
        },
        
        'Asparagine': {
            'folder': 'Asparagine-Reference',
            # Asparagine: H2N-C(=O)-CH2-CH(NH2)-COOH
            # β-CH2 at ~2.84 ppm (2H) - DIASTEREOTOPIC: splits into 2 peaks (~2.75, ~2.85)
            # α-CH at ~3.98 ppm (1H)
            # NOTE: β-CH2 peaks shift with concentration (2.71-2.85 ppm), widen region to capture
            'regions': [
                (2.65, 3.00),   # β-CH2 (2 peaks, diastereotopic) - widened for concentration shift
                (3.85, 4.10),   # α-CH
            ],
            'region_peaks': [
                [2.75, 2.85],   # β-CH2 (diastereotopic, 2 peaks)
                [3.98],
            ],
            'region_protons': [2, 1],
            'region_names': ['β-CH2 (2 peaks)', 'α-CH'],
            'ref_conc': 15.0,
            'files': {10: 15.0, 20: 7.5, 30: 3.75, 40: 1.875,
                     50: 0.9375, 60: 0.46875, 70: 0.234375}
        },
        
        'Phenylalanine': {
            'folder': 'Phenylalanine-Reference',
            # Phenylalanine: Phenyl-CH2-CH(NH2)-COOH
            # Aromatic ring: H-2,6 (ortho) at ~7.43 ppm, H-3,5 (meta) at ~7.35 ppm, H-4 (para) at ~7.28 ppm
            # Total 5 aromatic protons - appear as complex multiplet
            # Also has β-CH2 at ~3.1 ppm (2H, doublet of doublets)
            'regions': [
                (7.20, 7.55),   # Aromatic ring (5H, complex multiplet)
                (3.00, 3.25),   # β-CH2 (2H)
            ],
            'region_peaks': [
                [7.335, 7.435],  # Aromatic envelope (2 main peaks)
                [3.12],          # β-CH2
            ],
            'region_protons': [5, 2],
            'region_names': ['Ar-H (5H)', 'β-CH2'],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.400125, 30: 1.200063, 40: 0.600031,
                     50: 0.300016, 60: 0.150008, 70: 0.075004}
        },
        
        'Tyrosine': {
            'folder': 'Tyrosine-Reference',
            # Tyrosine: p-OH-Ph-CH2-CH(NH2)-COOH
            # Aromatic protons at 6.91 and 7.20 ppm (4H total)
            # β-CH2 at ~3.05 ppm (2H)
            'regions': [
                (6.80, 7.30),   # Aromatic AA'BB' system (4H)
                (2.90, 3.20),   # β-CH2 (2H)
            ],
            'region_peaks': [
                [6.91, 7.20],
                [3.05],
            ],
            'region_protons': [4, 2],
            'region_names': ['Ar-H (4H)', 'β-CH2'],
            'ref_conc': 2.0,
            'files': {10: 2.0, 20: 1.087256, 30: 0.543628, 40: 0.271814,
                     50: 0.135907, 60: 0.067954, 70: 0.033977}
        },
        
        'Methionine': {
            'folder': 'Methionine-Reference',
            # Methionine: CH3-S-CH2-CH2-CH(NH2)-COOH
            # CH3 at ~2.15 ppm (3H), γ-CH2 at ~2.65 ppm (2H), β-CH2 at ~2.15 ppm (2H)
            # Note: CH3 and β-CH2 overlap at ~2.15 ppm
            'regions': [
                (2.05, 2.25),   # CH3 + β-CH2 (5H total)
                (2.55, 2.75),   # γ-CH2 (2H)
            ],
            'region_peaks': [
                [2.15],
                [2.65],
            ],
            'region_protons': [5, 2],
            'region_names': ['CH3+β-CH2', 'γ-CH2'],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.598861, 30: 1.299431, 40: 0.649715,
                     50: 0.324858, 60: 0.162429, 70: 0.081214}
        },
    }
    
    return metabolites


def fit_region(x, y, peak_centers, n_peaks, ref_sigma=None):
    """
    Fit Lorentzian model to spectral region.
    
    Parameters:
    -----------
    x, y : arrays
        PPM and intensity data
    peak_centers : list
        Initial peak center positions
    n_peaks : int
        Number of peaks to fit
    ref_sigma : float, optional
        Reference sigma value (if None, varies sigma)
    
    Returns:
    --------
    result : lmfit ModelResult
        Fit result object
    model : lmfit Model
        Fitted model
    """
    if n_peaks == 1:
        model = LorentzianModel() + ConstantModel()
        pars = model.make_params()
        pars['amplitude'].set(value=np.max(y)*0.01, min=0)
        pars['center'].set(value=peak_centers[0], min=np.min(x), max=np.max(x))
        if ref_sigma:
            pars['sigma'].set(value=ref_sigma, vary=False)
        else:
            pars['sigma'].set(value=0.005, min=0.001, max=0.02)
        pars['c'].set(value=np.min(y))
    elif n_peaks == 2:
        model = LorentzianModel(prefix='p1_') + LorentzianModel(prefix='p2_') + ConstantModel()
        pars = model.make_params()
        pars['p1_amplitude'].set(value=np.max(y)*0.01, min=0)
        pars['p1_center'].set(value=peak_centers[0], min=np.min(x), max=np.max(x))
        pars['p2_amplitude'].set(value=np.max(y)*0.01, min=0)
        pars['p2_center'].set(value=peak_centers[1], min=np.min(x), max=np.max(x))
        if ref_sigma:
            pars['p1_sigma'].set(value=ref_sigma, vary=False)
            pars['p2_sigma'].set(value=ref_sigma, vary=False)
        else:
            pars['p1_sigma'].set(value=0.005, min=0.001, max=0.02)
            pars['p2_sigma'].set(value=0.005, min=0.001, max=0.02)
        pars['c'].set(value=np.min(y))
    elif n_peaks == 3:
        # Support for 3 peaks (e.g., Arginine with split γ-CH2)
        model = (LorentzianModel(prefix='p1_') + LorentzianModel(prefix='p2_') + 
                 LorentzianModel(prefix='p3_') + ConstantModel())
        pars = model.make_params()
        for i, prefix in enumerate(['p1_', 'p2_', 'p3_'], 1):
            pars[f'{prefix}amplitude'].set(value=np.max(y)*0.01, min=0)
            pars[f'{prefix}center'].set(value=peak_centers[i-1], min=np.min(x), max=np.max(x))
            if ref_sigma:
                pars[f'{prefix}sigma'].set(value=ref_sigma, vary=False)
            else:
                pars[f'{prefix}sigma'].set(value=0.005, min=0.001, max=0.02)
        pars['c'].set(value=np.min(y))
    else:
        raise ValueError(f"Unsupported number of peaks: {n_peaks}")
    
    result = model.fit(y, pars, x=x)
    return result, model


def quantify_metabolite_v3(met_name, met_info, base_dir, output_dir):
    """
    Quantify a single metabolite using multi-region fitting.
    
    Returns:
    --------
    summary : dict
        Contains results for each region and combined weighted result
    """
    folder = met_info['folder']
    regions = met_info['regions']
    region_peaks = met_info['region_peaks']
    region_protons = met_info['region_protons']
    region_names = met_info['region_names']
    files = met_info['files']
    
    folder_path = os.path.join(base_dir, folder)
    if not os.path.exists(folder_path):
        print(f"  Warning: Folder {folder_path} not found")
        return None
    
    # Check available files
    available_files = {}
    for fileno, conc in files.items():
        filepath = os.path.join(folder_path, f"{fileno}.dx")
        if os.path.exists(filepath):
            available_files[fileno] = conc
    
    if len(available_files) < 2:
        print(f"  Warning: Not enough files for {met_name}")
        return None
    
    print(f"  Processing {met_name}: {len(regions)} regions, {len(available_files)} files")
    
    # Exclude file 10 from analysis
    if 10 in available_files:
        del available_files[10]
    
    # Use file 20 as reference instead of file 10
    if 20 not in available_files:
        print(f"  Warning: File 20 not found for {met_name}")
        return None
    
    ref_fileno = 20
    ref_conc = available_files[ref_fileno]
    
    # Read reference
    try:
        ppm_ref, spec_ref = read_and_process(os.path.join(folder_path, f"{ref_fileno}.dx"))
    except Exception as e:
        print(f"  Error reading reference: {e}")
        return None
    
    tsp_ref = find_tsp_peak(ppm_ref, spec_ref)
    ppm_ref_corr = ppm_ref - tsp_ref
    tsp_area_ref = integrate_peak(ppm_ref_corr, spec_ref, (-0.2, 0.2))
    
    # Fit reference for each region
    ref_fits = []
    for i, (region, peaks) in enumerate(zip(regions, region_peaks)):
        mask = (ppm_ref_corr >= region[0]) & (ppm_ref_corr <= region[1])
        x_ref = ppm_ref_corr[mask]
        y_ref = spec_ref[mask] / tsp_area_ref
        
        if len(x_ref) == 0:
            ref_fits.append(None)
            continue
        
        n_peaks = len(peaks)
        result_ref, model = fit_region(x_ref, y_ref, peaks, n_peaks)
        
        if n_peaks == 1:
            ref_amp = result_ref.params['amplitude'].value
            ref_sigma = result_ref.params['sigma'].value
        elif n_peaks == 2:
            ref_amp = result_ref.params['p1_amplitude'].value + result_ref.params['p2_amplitude'].value
            ref_sigma = (result_ref.params['p1_sigma'].value + result_ref.params['p2_sigma'].value) / 2
        elif n_peaks == 3:
            ref_amp = (result_ref.params['p1_amplitude'].value + result_ref.params['p2_amplitude'].value + 
                      result_ref.params['p3_amplitude'].value)
            ref_sigma = (result_ref.params['p1_sigma'].value + result_ref.params['p2_sigma'].value + 
                        result_ref.params['p3_sigma'].value) / 3
        
        ref_fits.append({
            'model': model,
            'amplitude': ref_amp,
            'sigma': ref_sigma,
            'result': result_ref,
            'n_peaks': n_peaks,
            'peaks': peaks
        })
    
    # Quantify all samples
    region_results = [[] for _ in regions]  # Results per region
    
    for fileno, true_conc in sorted(available_files.items()):
        if fileno == ref_fileno:
            # Reference file - concentration is known
            for i, region_res in enumerate(region_results):
                region_res.append({
                    'fileno': fileno,
                    'true': true_conc,
                    'calc': ref_conc,
                    'recovery': 100.0,
                    'scale': 1.0,
                    'scale_tsp': 1.0,
                    'r2': ref_fits[i]['result'].rsquared if ref_fits[i] else 0
                })
            continue
        
        # Read sample
        try:
            ppm_samp, spec_samp = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
        except:
            continue
        
        tsp_samp = find_tsp_peak(ppm_samp, spec_samp)
        ppm_samp_corr = ppm_samp - tsp_samp
        tsp_area_samp = integrate_peak(ppm_samp_corr, spec_samp, (-0.2, 0.2))
        scale_tsp = tsp_area_samp / tsp_area_ref
        
        # Fit each region
        for i, (region, ref_fit) in enumerate(zip(regions, ref_fits)):
            if ref_fit is None:
                continue
            
            mask = (ppm_samp_corr >= region[0]) & (ppm_samp_corr <= region[1])
            x_samp = ppm_samp_corr[mask]
            y_samp = spec_samp[mask] / tsp_area_samp
            
            if len(x_samp) == 0:
                continue
            
            try:
                # Detect actual peak positions
                detected_centers = detect_peaks_in_region(
                    ppm_samp_corr, spec_samp / tsp_area_samp, region, 
                    ref_fit['peaks'], ref_fit['n_peaks']
                )
                
                # Fit with detected centers
                result_samp, _ = fit_region(x_samp, y_samp, detected_centers, 
                                           ref_fit['n_peaks'], ref_fit['sigma'])
                
                # Calculate scale
                if ref_fit['n_peaks'] == 1:
                    samp_amp = result_samp.params['amplitude'].value
                elif ref_fit['n_peaks'] == 2:
                    samp_amp = (result_samp.params['p1_amplitude'].value + 
                               result_samp.params['p2_amplitude'].value)
                elif ref_fit['n_peaks'] == 3:
                    samp_amp = (result_samp.params['p1_amplitude'].value + 
                               result_samp.params['p2_amplitude'].value +
                               result_samp.params['p3_amplitude'].value)
                
                scale = samp_amp / ref_fit['amplitude'] if ref_fit['amplitude'] > 1e-10 else 0
                calc_conc = ref_conc * scale
                recovery = 100 * calc_conc / true_conc if true_conc > 0 else 0
                
                region_results[i].append({
                    'fileno': fileno,
                    'true': true_conc,
                    'calc': calc_conc,
                    'recovery': recovery,
                    'scale': scale,
                    'scale_tsp': scale_tsp,
                    'r2': result_samp.rsquared
                })
            except Exception as e:
                print(f"    Error fitting region {i} file {fileno}: {e}")
                continue
    
    # Calculate proton-weighted average concentration
    combined_results = []
    for idx, fileno in enumerate(sorted(available_files.keys())):
        true_conc = available_files[fileno]
        
        # Collect valid region results for this file
        weighted_sum = 0
        total_protons = 0
        region_calcs = []
        
        for i, region_res in enumerate(region_results):
            # Find this file in region results
            file_result = next((r for r in region_res if r['fileno'] == fileno), None)
            if file_result and file_result['calc'] > 0:
                protons = region_protons[i]
                weighted_sum += file_result['calc'] * protons
                total_protons += protons
                region_calcs.append({
                    'region': region_names[i],
                    'calc': file_result['calc'],
                    'protons': protons,
                    'recovery': file_result['recovery']
                })
        
        if total_protons > 0:
            combined_calc = weighted_sum / total_protons
            combined_recovery = 100 * combined_calc / true_conc if true_conc > 0 else 0
        else:
            combined_calc = 0
            combined_recovery = 0
        
        combined_results.append({
            'fileno': fileno,
            'true': true_conc,
            'calc': combined_calc,
            'recovery': combined_recovery,
            'region_details': region_calcs
        })
    
    # Calculate statistics
    recoveries = [r['recovery'] for r in combined_results]
    mean_recovery = np.mean(recoveries)
    std_recovery = np.std(recoveries)
    
    true_vals = [r['true'] for r in combined_results]
    calc_vals = [r['calc'] for r in combined_results]
    
    if len(true_vals) >= 2:
        slope, intercept = np.polyfit(true_vals, calc_vals, 1)
        r_squared = np.corrcoef(true_vals, calc_vals)[0, 1]**2
    else:
        slope, intercept, r_squared = 0, 0, 0
    
    # Generate plot with dynamic layout based on number of regions
    n_regions = len(regions)
    # Layout: 3 rows, but row 2 shows all regions side by side
    # If more than 2 regions, extend figure width
    fig_width = 16 if n_regions <= 2 else 8 + 4 * n_regions
    fig_height = 12 if n_regions <= 2 else 10 + 2 * n_regions
    
    fig = plt.figure(figsize=(fig_width, fig_height))
    gs = fig.add_gridspec(3, max(3, n_regions + 1), hspace=0.3, wspace=0.3)
    
    # Row 1: Calibration curve and region comparison
    ax1 = fig.add_subplot(gs[0, :2])
    ax1.plot(true_vals, calc_vals, 'bo', markersize=10, label='Weighted Average')
    ax1.plot(true_vals, true_vals, 'k--', label='Ideal (y=x)')
    if r_squared > 0:
        ax1.plot(true_vals, np.polyval([slope, intercept], true_vals), 'r-',
                 label=f'Fit: y={slope:.3f}x+{intercept:.3f}, R²={r_squared:.4f}')
    ax1.set_xlabel('True Concentration (mM)', fontsize=12)
    ax1.set_ylabel('Calculated Concentration (mM)', fontsize=12)
    ax1.set_title(f'{met_name} Multi-Region Quantification', fontsize=14)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Region comparison table with ranges
    ax2 = fig.add_subplot(gs[0, 2])
    ax2.axis('off')
    ax2.set_title(f'Region Comparison (File {ref_fileno})', fontsize=11, fontweight='bold')
    
    table_data = []
    for i, (name, protons, region) in enumerate(zip(region_names, region_protons, regions)):
        region_range = f'{region[0]:.2f}-{region[1]:.2f}'
        table_data.append([name, region_range, str(protons)])
    
    table = ax2.table(
        cellText=table_data,
        colLabels=['Region', 'Range (ppm)', 'Protons'],
        loc='center',
        cellLoc='center',
        colColours=['#4472C4']*3,
        colWidths=[0.4, 0.35, 0.25]
    )
    table.auto_set_font_size(False)
    table.set_fontsize(8)
    table.scale(1.2, 1.5)
    for i in range(3):
        table[(0, i)].set_text_props(color='white', fontweight='bold')
    
    # Row 2: Show ALL regions
    region_axes = []
    for i in range(n_regions):
        if i < n_regions:
            ax = fig.add_subplot(gs[1, i])
            region_axes.append(ax)
            
            colors = plt.cm.viridis(np.linspace(0, 1, len(combined_results)))
            for j, (r, color) in enumerate(zip(combined_results, colors)):
                fileno = r['fileno']
                try:
                    ppm, spec = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
                    tsp = find_tsp_peak(ppm, spec)
                    ppm_corr = ppm - tsp
                    
                    region = regions[i]
                    mask = (ppm_corr >= region[0]) & (ppm_corr <= region[1])
                    if np.sum(mask) > 0:
                        ax.plot(ppm_corr[mask], spec[mask], color=color,
                                label=f'{fileno}', linewidth=1.5, alpha=0.8)
                except:
                    pass
            
            ax.set_xlabel('Chemical Shift (ppm)', fontsize=10)
            ax.set_ylabel('Raw Intensity', fontsize=10)
            ax.set_title(f'Region {i+1}: {region_names[i]}\n({regions[i][0]:.2f}-{regions[i][1]:.2f} ppm)', 
                        fontsize=10)
            ax.set_xlim(regions[i][1], regions[i][0])
            ax.grid(True, alpha=0.3)
            if i == 0:
                ax.legend(fontsize=5, loc='upper right')
    
    # TSP reference (last column of row 2)
    ax_tsp = fig.add_subplot(gs[1, min(n_regions, max(3, n_regions + 1) - 1)])
    colors = plt.cm.viridis(np.linspace(0, 1, len(combined_results)))
    for j, (r, color) in enumerate(zip(combined_results, colors)):
        fileno = r['fileno']
        try:
            ppm, spec = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            mask_tsp = (ppm_corr >= -0.2) & (ppm_corr <= 0.2)
            if np.sum(mask_tsp) > 0:
                ax_tsp.plot(ppm_corr[mask_tsp], spec[mask_tsp], color=color,
                        label=f'{fileno}', linewidth=1.5)
        except:
            pass
    ax_tsp.set_xlabel('Chemical Shift (ppm)', fontsize=10)
    ax_tsp.set_ylabel('Raw Intensity', fontsize=10)
    ax_tsp.set_title('TSP Reference', fontsize=10)
    ax_tsp.set_xlim(0.2, -0.2)
    ax_tsp.grid(True, alpha=0.3)
    
    # Row 3: Fits for representative samples (first region)
    # Plot file 20, file 30, and file 40
    target_files = [20, 30, 40]
    sample_indices = []
    for target in target_files:
        found = False
        for idx, r in enumerate(combined_results):
            if r['fileno'] == target:
                sample_indices.append(idx)
                found = True
                break
        if not found and len(combined_results) > 0:
            # If target not found, use first available
            sample_indices.append(0)
    
    # Ensure we have at most 3 plots
    sample_indices = sample_indices[:3]
    
    for plot_idx, result_idx in enumerate(sample_indices):
        if plot_idx >= 3:
            break
        ax = fig.add_subplot(gs[2, plot_idx])
        r = combined_results[result_idx]
        fileno = r['fileno']
        
        # Show primary region fit
        try:
            ppm, spec = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            
            region = regions[0]
            ref_fit = ref_fits[0]
            
            if ref_fit is None:
                ax.text(0.5, 0.5, 'No ref fit', ha='center', va='center', transform=ax.transAxes)
                continue
            
            # Step 1: Detect peaks in the wide region to find actual peak positions
            detected = detect_peaks_in_region(ppm_corr, spec/tsp_area_ref, region,
                                              ref_fit['peaks'], ref_fit['n_peaks'])
            
            # Step 2: Calculate dynamic range based on detected peaks
            # Use expected peaks as fallback if detection fails
            if len(detected) == 0:
                detected = ref_fit['peaks']
            
            # Ensure we have the right number of peaks
            if len(detected) != ref_fit['n_peaks']:
                # Pad with expected positions or truncate
                if len(detected) < ref_fit['n_peaks']:
                    detected = list(detected) + ref_fit['peaks'][len(detected):]
                else:
                    detected = detected[:ref_fit['n_peaks']]
            
            # Calculate dynamic range: peaks ± 0.20 ppm padding
            dynamic_min = min(detected) - 0.20
            dynamic_max = max(detected) + 0.20
            
            # Step 3: Extract data using dynamic range (data = fit = plot range)
            mask = (ppm_corr >= dynamic_min) & (ppm_corr <= dynamic_max)
            x = ppm_corr[mask]
            y = spec[mask] / tsp_area_ref
            
            if len(x) > 0:
                result, _ = fit_region(x, y, detected, ref_fit['n_peaks'], ref_fit['sigma'])
                
                ax.plot(x, y, 'b.', markersize=3, label='Data')
                ax.plot(x, result.best_fit, 'r-', linewidth=1.5, label='Fit')
                ax.set_title(f'File {fileno}: True={r["true"]:.2f} mM\n'
                           f'Calc={r["calc"]:.2f} mM ({r["recovery"]:.1f}%)', fontsize=10)
                ax.set_xlabel('ppm')
                ax.set_ylabel('Intensity (norm)')
                ax.legend(fontsize=8)
                
                # Step 4: Plot xlim matches data range exactly
                ax.set_xlim(dynamic_max, dynamic_min)
        except Exception as e:
            ax.text(0.5, 0.5, f'Error: {str(e)[:50]}', ha='center', va='center', transform=ax.transAxes)
    
    output_path = os.path.join(output_dir, f'{met_name}_v3_ref20_quantification.png')
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    summary = {
        'name': met_name,
        'ref_conc': ref_conc,
        'n_regions': len(regions),
        'region_names': region_names,
        'region_protons': region_protons,
        'slope': slope,
        'intercept': intercept,
        'r_squared': r_squared,
        'mean_recovery': mean_recovery,
        'std_recovery': std_recovery,
        'num_samples': len(combined_results),
        'results': combined_results,
        'region_details': region_results
    }
    
    return summary


def main():
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX"
    output_dir = "quantification_results"
    
    os.makedirs(output_dir, exist_ok=True)
    
    print("="*100)
    print("ABSOLUTE QUANTIFICATION V3 - PHYSICS-BASED MULTI-REGION FITTING")
    print("="*100)
    print()
    print("Key features:")
    print("  • Multi-region fitting: Uses ALL distinct chemical shift regions")
    print("  • Proton-weighted averaging: Combines regions by proton count")
    print("  • Physics-based peak counting: Based on molecular structure")
    print("  • Internal consistency check: Validates across regions")
    print()
    
    metabolites = get_metabolite_info_v3()
    all_summaries = []
    
    for met_name, met_info in metabolites.items():
        print(f"\nProcessing {met_name}...")
        summary = quantify_metabolite_v3(met_name, met_info, base_dir, output_dir)
        if summary:
            all_summaries.append(summary)
            print(f"  ✓ {met_name}: {summary['n_regions']} regions, "
                  f"R²={summary['r_squared']:.4f}, Recovery={summary['mean_recovery']:.1f}%")
        else:
            print(f"  ✗ {met_name}: Failed")
    
    # Print overall summary
    print()
    print("="*100)
    print("OVERALL SUMMARY V3")
    print("="*100)
    print(f"{'Metabolite':<20} {'Regions':<8} {'R²':<10} {'Mean Rec (%)':<15} {'SD (%)':<10} {'N':<5}")
    print("-"*100)
    for s in all_summaries:
        print(f"{s['name']:<20} {s['n_regions']:<8} {s['r_squared']:<10.4f} "
              f"{s['mean_recovery']:<15.1f} {s['std_recovery']:<10.1f} {s['num_samples']:<5}")
    print("="*100)
    
    # Save detailed results
    with open(os.path.join(output_dir, 'quantification_summary_v3.csv'), 'w') as f:
        f.write("Metabolite,File,True_Conc_mM,Calc_Conc_mM,Recovery_pct,R2\n")
        for s in all_summaries:
            for r in s['results']:
                f.write(f"{s['name']},{r['fileno']},{r['true']:.6f},"
                       f"{r['calc']:.6f},{r['recovery']:.2f},{s['r_squared']:.4f}\n")
    
    print(f"\nResults saved to {output_dir}/")
    print(f"  - Individual PNG plots (*_v3_quantification.png)")
    print(f"  - quantification_summary_v3.csv")


if __name__ == '__main__':
    main()