#python master_feature_ode_solver.py music.mp3 speech.mp3 audio audio
#python master_feature_ode_solver.py sunset.mp4 rain.mp4 video video

import numpy as np
import torch
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import librosa
import cv2
import os

# ------------------------------------------------------------
# 1. Master Feature Signal Extraction from Media
# ------------------------------------------------------------
def extract_master_feature_audio(mp3_path, target_length=1000):
    """
    Load MP3 and extract a 1D master feature signal (e.g., amplitude envelope).
    Returns a numpy array of shape (target_length,).
    """
    y, sr = librosa.load(mp3_path, sr=None)
    # Compute amplitude envelope by taking RMS in frames
    hop_length = max(1, len(y) // target_length)
    envelope = librosa.feature.rms(y=y, hop_length=hop_length, frame_length=hop_length*2)[0]
    # Resample to exactly target_length
    if len(envelope) > target_length:
        envelope = envelope[:target_length]
    else:
        envelope = np.pad(envelope, (0, target_length - len(envelope)), mode='constant')
    # Normalize to [0,1] range
    envelope = (envelope - envelope.min()) / (envelope.max() - envelope.min() + 1e-8)
    return envelope

def extract_master_feature_video(mp4_path, target_length=1000):
    """
    Load MP4 and extract a 1D master feature signal as average frame brightness.
    Returns a numpy array of shape (target_length,).
    """
    cap = cv2.VideoCapture(mp4_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    duration = frame_count / fps if fps > 0 else 1.0
    brightness_series = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        brightness_series.append(np.mean(gray))
    cap.release()
    brightness_series = np.array(brightness_series)
    # Resample to target_length
    indices = np.linspace(0, len(brightness_series)-1, target_length).astype(int)
    resampled = brightness_series[indices]
    # Normalize to [0,1]
    resampled = (resampled - resampled.min()) / (resampled.max() - resampled.min() + 1e-8)
    return resampled

# ------------------------------------------------------------
# 2. ODE Definition using the Master Feature Signal
# ------------------------------------------------------------
def ode_model(y, t, phi_interp):
    """
    Undetermined ODE: dy/dt = phi(t) * y
    Without phi(t) this equation has infinite solutions.
    phi_interp is a function that returns the master feature value at time t.
    """
    phi_t = phi_interp(t)
    return phi_t * y

def solve_ode_with_master_feature(phi_signal, t_eval):
    """
    Solves dy/dt = phi(t) * y, y(0)=1.
    phi_signal: 1D array of master feature values at equally spaced times.
    t_eval: times where solution is desired (must be within [0,1]).
    """
    from scipy.interpolate import interp1d
    # Normalize time axis to [0,1] for the signal
    t_signal = np.linspace(0, 1, len(phi_signal))
    phi_interp = interp1d(t_signal, phi_signal, kind='linear', fill_value=0, bounds_error=False)

    # Initial condition
    y0 = 1.0
    sol = odeint(ode_model, y0, t_eval, args=(phi_interp,))
    return sol.flatten()

# ------------------------------------------------------------
# 3. Comparison of Two Different Media Sources
# ------------------------------------------------------------
def compare_media_sources(file1, file2, file1_type='audio', file2_type='audio', num_points=200):
    """
    Extracts master features from two files, solves the same ODE, and plots.
    file_type: 'audio' or 'video'
    """
    # Extract master feature signals
    if file1_type == 'audio':
        phi1 = extract_master_feature_audio(file1, target_length=num_points)
    else:
        phi1 = extract_master_feature_video(file1, target_length=num_points)

    if file2_type == 'audio':
        phi2 = extract_master_feature_audio(file2, target_length=num_points)
    else:
        phi2 = extract_master_feature_video(file2, target_length=num_points)

    # Time grid for ODE solution (from 0 to 1)
    t_eval = np.linspace(0, 1, num_points)

    # Solve ODE using each master feature
    sol1 = solve_ode_with_master_feature(phi1, t_eval)
    sol2 = solve_ode_with_master_feature(phi2, t_eval)

    # Plot results
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(t_eval, phi1, label=f'Master feature 1\n({os.path.basename(file1)})')
    plt.plot(t_eval, phi2, label=f'Master feature 2\n({os.path.basename(file2)})')
    plt.title('Master Feature Signals (normalized)')
    plt.xlabel('Time (normalized)')
    plt.ylabel('Amplitude')
    plt.legend()
    plt.grid(alpha=0.3)

    plt.subplot(1, 2, 2)
    plt.plot(t_eval, sol1, label=f'ODE solution from {os.path.basename(file1)}')
    plt.plot(t_eval, sol2, label=f'ODE solution from {os.path.basename(file2)}')
    plt.title('Solution of dy/dt = φ(t)·y , y(0)=1')
    plt.xlabel('Time')
    plt.ylabel('y(t)')
    plt.legend()
    plt.grid(alpha=0.3)

    plt.tight_layout()
    plt.show()

    # Compute and print similarity metrics
    mse = np.mean((sol1 - sol2)**2)
    corr = np.corrcoef(sol1, sol2)[0,1]
    print(f"Comparison between {file1} and {file2}:")
    print(f"  MSE between solutions: {mse:.6f}")
    print(f"  Correlation coefficient: {corr:.6f}")
    print(sol1)
    print(sol2)

    return sol1, sol2, phi1, phi2

# ------------------------------------------------------------
# 4. Example Usage
# ------------------------------------------------------------
if __name__ == "__main__":
    # Replace with your actual file paths
    # audio1 = "path/to/song.mp3"
    # audio2 = "path/to/speech.mp3"
    # video1 = "path/to/scene.mp4"
    # video2 = "path/to/another.mp4"

    # For demonstration, we create synthetic files if real ones are missing
    # (In practice, use your own media files)
    print("Master Feature Signal ODE Solver")
    #print("Please provide real MP3/MP4 file paths, or the script will use synthetic signals.")
    
    # Example with a sine wave if no files given
    import sys
    if len(sys.argv) >= 3:
        # Command line arguments: python script.py file1 file2 [type1] [type2]
        f1 = sys.argv[1]
        f2 = sys.argv[2]
        t1 = sys.argv[3] if len(sys.argv) > 3 else 'audio'
        t2 = sys.argv[4] if len(sys.argv) > 4 else 'audio'
        compare_media_sources(f1, f2, t1, t2)
    else:
        # Create synthetic example: two different artificial signals
        t_synth = np.linspace(0, 1, 200)
        phi_synth1 = 0.5 + 0.5 * np.sin(2 * np.pi * 3 * t_synth)      # oscillatory
        phi_synth2 = 0.2 + 0.8 * np.exp(-5 * t_synth)                  # decay

        # Solve ODE with these "master features"
        sol1 = solve_ode_with_master_feature(phi_synth1, t_synth)
        sol2 = solve_ode_with_master_feature(phi_synth2, t_synth)

        plt.figure(figsize=(12,5))
        plt.subplot(1,2,1)
        plt.plot(t_synth, phi_synth1, label='Synthetic master 1 (sinusoid)')
        plt.plot(t_synth, phi_synth2, label='Synthetic master 2 (exponential decay)')
        plt.title('Synthetic Master Features')
        plt.legend()
        plt.subplot(1,2,2)
        plt.plot(t_synth, sol1, label='Solution from sin master')
        plt.plot(t_synth, sol2, label='Solution from exp master')
        plt.title('ODE Solutions')
        plt.legend()
        plt.show()
        print("\nDemo with synthetic signals. Replace with real media files for true experiments.")