from flask import Flask, request, jsonify, make_response
from flask_cors import CORS
import numpy as np
import io
import os
import tensorflow as tf
from scipy.io import wavfile
from scipy.fft import fft

app = Flask(__name__)
CORS(app)

# Load VAD model
print("Chargement du modele VAD (CRNN)...")
model_vad = tf.keras.models.load_model("CRNN_model_final.h5", compile=False)
print(f"VAD Input: {model_vad.input_shape}, Output: {model_vad.output_shape}")

# Load DOA model
print("Chargement du modele DOA...")
model_doa = tf.keras.models.load_model("model_keras.h5", compile=False)
print(f"DOA Input: {model_doa.input_shape}, Output: {model_doa.output_shape}")

SILENCE_CLASSES = [13, 14]
VAD_SEGMENT_SIZE = 1024

# DOA parameters (matching MATLAB code exactly)
M = 4                    # nombre de microphones
D_INTER = 0.02           # espacement 2cm
WIN_LEN = 1024           # taille fenetre STFT
HOP = 512                # hop size
N_SNAPSHOTS = 10         # snapshots pour moyennage
FS = 16000               # frequence echantillonnage
N_BLOCKS = 34            # nombre de trames

# Frequency bins 500-4000 Hz
f_axis = np.arange(0, WIN_LEN // 2 + 1) * (FS / WIN_LEN)
v_bins = np.where((f_axis >= 500) & (f_axis <= 4000))[0]
N_BINS = len(v_bins)  # should be ~225

# Co-array parameters for ULA with 4 mics
# For ULA: differences are -3, -2, -1, 0, 1, 2, 3 (s=7)
S = 2 * M - 1  # = 7 for 4 mics
ZERO_IDX = M - 1  # = 3 (0-indexed: position of 0 in diffs)
M_V = M  # = 4
LEN_FEAT = S - 4  # = 3 (indices 4 to 6, i.e., indices 5:end in MATLAB 1-indexed)

def get_coarray_indices(m):
    """Generate co-array index mapping for ULA"""
    # For ULA with m microphones, differences range from -(m-1) to (m-1)
    # index_map[d] = list of (i,j) pairs where j-i = d
    index_map = {}
    for d in range(-(m-1), m):
        pairs = []
        for i in range(m):
            for j in range(m):
                if j - i == d:
                    pairs.append((i, j))
        index_map[d] = pairs
    return index_map

INDEX_MAP = get_coarray_indices(M)

def predict_vad(audio_segment):
    if len(audio_segment) < VAD_SEGMENT_SIZE:
        audio_segment = np.pad(audio_segment, (0, VAD_SEGMENT_SIZE - len(audio_segment)), mode='constant')
    elif len(audio_segment) > VAD_SEGMENT_SIZE:
        start = (len(audio_segment) - VAD_SEGMENT_SIZE) // 2
        audio_segment = audio_segment[start:start + VAD_SEGMENT_SIZE]
    
    max_val = np.max(np.abs(audio_segment))
    if max_val > 0:
        audio_segment = audio_segment / max_val
    
    features = audio_segment.reshape(1, VAD_SEGMENT_SIZE, 1).astype(np.float32)
    prediction = model_vad.predict(features, verbose=0)
    predicted_class = int(np.argmax(prediction[0]))
    confidence = float(np.max(prediction[0]))
    label = "silence" if predicted_class in SILENCE_CLASSES else "voice"
    
    return label, confidence, predicted_class

def extract_doa_features(au_data):
    """
    Extract DOA features exactly as in MATLAB code.
    au_data: (n_samples, 4) - 4 channel audio
    Returns: (34, 225, 3) features
    """
    n_samples = au_data.shape[0]
    
    # Minimum length for 34 blocks
    min_len = (N_BLOCKS * N_SNAPSHOTS * HOP) + WIN_LEN
    
    # Repeat if too short
    if n_samples < min_len:
        rep = int(np.ceil(min_len / n_samples))
        au_data = np.tile(au_data, (rep, 1))
    
    # Take only min_len samples
    au_data = au_data[:min_len, :]
    
    # Initialize features array (34, 225, 3)
    frame_features = np.zeros((N_BLOCKS, N_BINS, LEN_FEAT))
    
    for k in range(N_BLOCKS):
        # Average covariance over snapshots
        R_avg = np.zeros((M, M, N_BINS), dtype=complex)
        
        for snap in range(N_SNAPSHOTS):
            idx_sample = k * N_SNAPSHOTS * HOP + snap * HOP
            segment = au_data[idx_sample:idx_sample + WIN_LEN, :]  # (1024, 4)
            
            # FFT for each channel
            X_f = fft(segment, n=WIN_LEN, axis=0)  # (1024, 4)
            X_f = X_f[v_bins, :].T  # (4, 225) - transposed to match MATLAB
            
            for b in range(N_BINS):
                vec_f = X_f[:, b]  # (4,) complex vector
                norm_sq = np.linalg.norm(vec_f) ** 2 + 1e-10
                R_avg[:, :, b] += np.outer(vec_f, np.conj(vec_f)) / norm_sq
        
        R_avg /= N_SNAPSHOTS
        
        # Extract co-array phases
        for b in range(N_BINS):
            R = R_avg[:, :, b]
            
            # Build co-array vector z
            z = np.zeros(S, dtype=complex)
            jj = 0
            for ii in range(-(M_V - 1), M_V):
                pairs = INDEX_MAP.get(ii, [])
                if pairs:
                    z[jj] = np.mean([R[i, j] for i, j in pairs])
                jj += 1
            
            # Extract phases from indices 4 to end (MATLAB: 5:s)
            phases = np.angle(z[4:])  # indices 4, 5, 6 -> 3 values
            frame_features[k, b, :] = phases
    
    return frame_features

def predict_doa(multichannel_audio):
    """
    Predict DOA angles using CNN model.
    multichannel_audio: (n_samples, 4) array
    Returns: list of 34 predicted angles
    """
    # Extract features (34, 225, 3)
    features = extract_doa_features(multichannel_audio)
    
    # Reshape to (1, 34, 675) - flatten last two dims
    features_flat = features.reshape(1, N_BLOCKS, -1).astype(np.float32)
    
    # Predict
    prediction = model_doa.predict(features_flat, verbose=0)
    angles = prediction[0].tolist()  # 34 angles
    
    return angles

@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', '*')
    response.headers.add('Access-Control-Allow-Methods', 'GET,POST,OPTIONS')
    return response

@app.route('/', methods=['GET'])
def home():
    return jsonify({'message': 'AcoustiTrack API', 'status': 'running', 'models': ['VAD-CRNN', 'DOA-CNN1D']})

@app.route('/api/health', methods=['GET', 'OPTIONS'])
def health():
    if request.method == 'OPTIONS':
        return make_response('', 204)
    return jsonify({'status': 'ok', 'model': 'CRNN VAD + CNN1D DOA'})

@app.route('/api/analyze', methods=['POST', 'OPTIONS'])
def analyze():
    if request.method == 'OPTIONS':
        return make_response('', 204)
    try:
        if 'audio' not in request.files:
            return jsonify({'error': 'Aucun fichier audio'}), 400
        
        audio_file = request.files['audio']
        audio_bytes = audio_file.read()
        
        # Read WAV file
        sr, data = wavfile.read(io.BytesIO(audio_bytes))
        
        # Convert to float
        if data.dtype == np.int16:
            data = data.astype(np.float32) / 32768.0
        elif data.dtype == np.int32:
            data = data.astype(np.float32) / 2147483648.0
        
        # Handle mono vs multi-channel
        if data.ndim == 1:
            # Mono - duplicate to 4 channels
            data_multi = np.tile(data.reshape(-1, 1), (1, 4))
            data_mono = data
        else:
            data_multi = data
            data_mono = np.mean(data, axis=1)
            # Ensure 4 channels
            if data_multi.shape[1] < 4:
                data_multi = np.hstack([data_multi, np.tile(data_multi[:, 0:1], (1, 4 - data_multi.shape[1]))])
            elif data_multi.shape[1] > 4:
                data_multi = data_multi[:, :4]
        
        n_channels = data_multi.shape[1]
        duration = len(data_mono) / sr
        
        # Resample to 16kHz if needed
        if sr != FS:
            from scipy.signal import resample
            new_len = int(len(data_mono) * FS / sr)
            data_mono = resample(data_mono, new_len)
            data_multi = np.column_stack([resample(data_multi[:, i], new_len) for i in range(4)])
            sr = FS
        
        # VAD prediction
        total_segments = len(data_mono) // VAD_SEGMENT_SIZE
        max_segments = min(total_segments, 20)
        if max_segments == 0:
            max_segments = 1
        step = max(1, total_segments // max_segments)
        
        vad_timeline = []
        voice_count = 0
        all_confidences = []
        
        for i in range(max_segments):
            segment_index = i * step
            start_sample = segment_index * VAD_SEGMENT_SIZE
            end_sample = start_sample + VAD_SEGMENT_SIZE
            
            if end_sample > len(data_mono):
                break
            
            segment = data_mono[start_sample:end_sample]
            label, confidence, _ = predict_vad(segment)
            
            if label == "voice":
                voice_count += 1
            all_confidences.append(confidence)
            
            vad_timeline.append({
                'start': float(round(start_sample / sr, 3)),
                'end': float(round(end_sample / sr, 3)),
                'label': label,
                'confidence': float(round(confidence, 3))
            })
        
        n_analyzed = len(vad_timeline)
        voice_ratio = float(voice_count / n_analyzed) if n_analyzed > 0 else 0.0
        vad_label = "voice" if voice_ratio > 0.5 else "silence"
        avg_vad_confidence = float(np.mean(all_confidences)) if all_confidences else 0.0
        
        # DOA prediction (uses 4-channel audio)
        doa_angles = predict_doa(data_multi)
        mean_angle = float(np.mean(doa_angles))
        std_angle = float(np.std(doa_angles))
        doa_confidence = float(max(0, 1 - std_angle / 25))
        
        # DOA trajectory
        doa_trajectory = []
        time_per_frame = duration / len(doa_angles)
        for i, angle in enumerate(doa_angles):
            doa_trajectory.append({
                'time': float(round(i * time_per_frame, 3)),
                'angle': float(round(angle, 1))
            })
        
        mean_y = float(np.mean(data_mono ** 2))
        snr = float(round(10 * np.log10(mean_y / 1e-10), 1)) if mean_y > 0 else 0.0
        
        return jsonify({
            'success': True,
            'metadata': {
                'channels': n_channels,
                'sampleRate': int(sr),
                'duration': float(round(duration, 2)),
                'samples': int(len(data_mono)),
                'segments_analyzed': n_analyzed
            },
            'vad': {
                'prediction': vad_label,
                'confidence': float(round(avg_vad_confidence, 3)),
                'timeline': vad_timeline
            },
            'doa': {
                'angle': float(round(mean_angle, 1)),
                'confidence': float(round(doa_confidence, 3)),
                'trajectory': doa_trajectory
            },
            'metrics': {
                'snr': snr,
                'voiceRatio': float(round(voice_ratio, 2)),
                'meanDoa': float(round(mean_angle, 1))
            }
        })
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))