/**
 * TensorFlow.js integration for neural network-based pitch detection
 * Handles model loading, inference, and tensor management
 */

import * as tf from '@tensorflow/tfjs';

/**
 * Model loading options
 */
export interface ModelLoadOptions {
  modelUrl?: string;
  backend?: 'cpu' | 'webgl' | 'wasm';
  warmup?: boolean;
}

/**
 * Inference result with predictions and confidence
 */
export interface InferenceResult {
  predictions: Float32Array;
  confidence: number;
  activations?: tf.Tensor;
}

/**
 * TensorFlow.js model manager
 */
export class TFJSModelManager {
  private model: tf.LayersModel | tf.GraphModel | null = null;
  private backend: string = 'webgl';
  private isInitialized: boolean = false;

  /**
   * Initialize TensorFlow.js with specified backend
   */
  async initialize(backend: 'cpu' | 'webgl' | 'wasm' = 'webgl'): Promise<void> {
    if (this.isInitialized) {
      return;
    }

    this.backend = backend;

    try {
      await tf.setBackend(backend);
      await tf.ready();
      this.isInitialized = true;
    } catch (error) {
      console.warn(`Failed to set backend ${backend}, falling back to CPU`, error);
      await tf.setBackend('cpu');
      await tf.ready();
      this.backend = 'cpu';
      this.isInitialized = true;
    }
  }

  /**
   * Load a model from URL
   */
  async loadModel(modelUrl: string, options: ModelLoadOptions = {}): Promise<void> {
    if (!this.isInitialized) {
      await this.initialize(options.backend);
    }

    try {
      // Try loading as LayersModel first
      this.model = await tf.loadLayersModel(modelUrl);
    } catch {
      // Fall back to GraphModel
      try {
        this.model = await tf.loadGraphModel(modelUrl);
      } catch (error) {
        throw new Error(`Failed to load model from ${modelUrl}: ${error}`);
      }
    }

    // Warmup the model with a dummy input
    if (options.warmup !== false) {
      await this.warmup();
    }
  }

  /**
   * Run inference on input data
   */
  async predict(input: Float32Array | tf.Tensor): Promise<tf.Tensor> {
    if (!this.model) {
      throw new Error('Model not loaded. Call loadModel() first.');
    }

    return tf.tidy(() => {
      // Convert input to tensor if needed
      let inputTensor: tf.Tensor;
      if (input instanceof Float32Array) {
        inputTensor = tf.tensor(Array.from(input));
      } else {
        inputTensor = input;
      }

      // Run prediction
      const output = this.model!.predict(inputTensor);
      
      // Handle both single and multiple outputs
      if (Array.isArray(output)) {
        return output[0] as tf.Tensor;
      }
      return output as tf.Tensor;
    });
  }

  /**
   * Run inference and return Float32Array
   */
  async predictArray(input: Float32Array, shape?: number[]): Promise<Float32Array> {
    const tensor = await this.predict(input);
    const data = await tensor.data();
    tensor.dispose();
    return new Float32Array(data);
  }

  /**
   * Batch prediction for multiple inputs
   */
  async predictBatch(inputs: Float32Array[]): Promise<Float32Array[]> {
    if (!this.model) {
      throw new Error('Model not loaded. Call loadModel() first.');
    }

    return tf.tidy(() => {
      // Stack inputs into a batch
      const tensors = inputs.map(input => tf.tensor(Array.from(input)));
      const batched = tf.stack(tensors);
      
      // Run prediction
      const output = this.model!.predict(batched) as tf.Tensor;
      
      // Unstack results
      const results: Float32Array[] = [];
      const outputArray = output.arraySync() as number[][];
      
      for (const row of outputArray) {
        results.push(new Float32Array(row));
      }
      
      return results;
    });
  }

  /**
   * Get model input shape
   */
  getInputShape(): number[] | null {
    if (!this.model) {
      return null;
    }

    const inputShape = this.model.inputs[0].shape;
    return inputShape as number[];
  }

  /**
   * Get model output shape
   */
  getOutputShape(): number[] | null {
    if (!this.model) {
      return null;
    }

    const outputShape = this.model.outputs[0].shape;
    return outputShape as number[];
  }

  /**
   * Warmup model with dummy input
   */
  private async warmup(): Promise<void> {
    if (!this.model) {
      return;
    }

    const inputShape = this.getInputShape();
    if (!inputShape) {
      return;
    }

    // Create dummy input matching expected shape
    const dummySize = inputShape.reduce((acc, dim) => {
      return acc * (dim === null || dim === -1 ? 1 : dim);
    }, 1);

    const dummyInput = new Float32Array(dummySize);
    
    try {
      const result = await this.predict(dummyInput);
      result.dispose();
    } catch (error) {
      console.warn('Model warmup failed:', error);
    }
  }

  /**
   * Get memory info
   */
  getMemoryInfo(): tf.MemoryInfo {
    return tf.memory();
  }

  /**
   * Dispose model and free memory
   */
  dispose(): void {
    if (this.model) {
      this.model.dispose();
      this.model = null;
    }
  }

  /**
   * Check if model is loaded
   */
  isLoaded(): boolean {
    return this.model !== null;
  }

  /**
   * Get current backend
   */
  getBackend(): string {
    return this.backend;
  }
}

/**
 * Tensor utilities for preprocessing and postprocessing
 */
export class TensorUtils {
  /**
   * Normalize audio to [-1, 1] range as tensor
   */
  static normalizeTensor(tensor: tf.Tensor): tf.Tensor {
    return tf.tidy(() => {
      const max = tensor.abs().max();
      return tensor.div(max);
    });
  }

  /**
   * Apply mean normalization to tensor
   */
  static meanNormalize(tensor: tf.Tensor): tf.Tensor {
    return tf.tidy(() => {
      const mean = tensor.mean();
      const std = tf.moments(tensor).variance.sqrt();
      return tensor.sub(mean).div(std.add(1e-8));
    });
  }

  /**
   * Convert Float32Array to batched tensor
   */
  static toBatchedTensor(data: Float32Array, batchSize: number = 1): tf.Tensor {
    return tf.tidy(() => {
      const tensor = tf.tensor(Array.from(data));
      return tensor.reshape([batchSize, -1]);
    });
  }

  /**
   * Apply softmax to get probability distribution
   */
  static softmax(tensor: tf.Tensor): tf.Tensor {
    return tf.tidy(() => {
      return tf.softmax(tensor);
    });
  }

  /**
   * Get argmax (index of maximum value)
   */
  static async argmax(tensor: tf.Tensor): Promise<number> {
    const argmaxTensor = tensor.argMax();
    const value = await argmaxTensor.data();
    argmaxTensor.dispose();
    return value[0];
  }

  /**
   * Get top-k indices and values
   */
  static async topK(tensor: tf.Tensor, k: number = 5): Promise<{ indices: number[]; values: number[] }> {
    const result = tf.topk(tensor, k);
    const indices = await result.indices.data();
    const values = await result.values.data();
    
    result.indices.dispose();
    result.values.dispose();
    
    return {
      indices: Array.from(indices),
      values: Array.from(values),
    };
  }

  /**
   * Dispose multiple tensors
   */
  static disposeTensors(...tensors: (tf.Tensor | null | undefined)[]): void {
    tensors.forEach(tensor => {
      if (tensor) {
        tensor.dispose();
      }
    });
  }
}

/**
 * Helper: Create model manager with auto-initialization
 */
export async function createModelManager(
  modelUrl: string,
  options: ModelLoadOptions = {}
): Promise<TFJSModelManager> {
  const manager = new TFJSModelManager();
  await manager.initialize(options.backend);
  await manager.loadModel(modelUrl, { ...options, warmup: options.warmup ?? true });
  return manager;
}

/**
 * Export TensorFlow.js for direct access if needed
 */
export { tf };