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

export interface ClassificationResult {
  label: string;
  confidence: number;
}

class CardModelService {
  private suitModel: tf.LayersModel | null = null;
  private valueModel: tf.LayersModel | null = null;

  private readonly SUIT_MODEL_PATH = '/models/suit_model/model.json';
  private readonly VALUE_MODEL_PATH = '/models/value_model/model.json';

  private readonly SUIT_LABELS = ['Schellen', 'Schilten', 'Eicheln', 'Rosen'];
  private readonly VALUE_LABELS = ['6', '7', '8', '9', '10', '11', '12', '13', '14'];

  async init(): Promise<void> {
    try {
      // Load models in parallel
      const [sModel, vModel] = await Promise.all([
        tf.loadLayersModel(this.SUIT_MODEL_PATH),
        tf.loadLayersModel(this.VALUE_MODEL_PATH),
      ]);
      this.suitModel = sModel;
      this.valueModel = vModel;
      console.log('Card detection models loaded successfully');
    } catch (error) {
      console.error('Failed to load card models:', error);
      throw error;
    }
  }

  async classifySuit(imageElement: HTMLCanvasElement | HTMLImageElement): Promise<ClassificationResult> {
    if (!this.suitModel) throw new Error('Suit model not initialized');

    return tf.tidy(() => {
      const tensor = this.preprocess(imageElement, 64);
      const prediction = this.suitModel!.predict(tensor) as tf.Tensor;
      const { index, confidence } = this.getTopK(prediction);
      
      return {
        label: this.SUIT_LABELS[index],
        confidence,
      };
    });
  }

  async classifyValue(imageElement: HTMLCanvasElement | HTMLImageElement): Promise<ClassificationResult> {
    if (!this.valueModel) throw new Error('Value model not initialized');

    return tf.tidy(() => {
      const tensor = this.preprocess(imageElement, 64);
      const prediction = this.valueModel!.predict(tensor) as tf.Tensor;
      const { index, confidence } = this.getTopK(prediction);
      
      return {
        label: this.VALUE_LABELS[index],
        confidence,
      };
    });
  }

  private preprocess(imageElement: HTMLCanvasElement | HTMLImageElement, size: number): tf.Tensor {
    return tf.browser.fromPixels(imageElement)
      .resizeBilinear([size, size])
      .expandDims(0)
      .div(255.0); // Normalize to [0, 1]
  }

  private getTopK(prediction: tf.Tensor): { index: number; confidence: number } {
    const data = prediction.dataSync();
    let maxIndex = 0;
    let maxVal = -1;

    for (let i = 0; i < data.length; i++) {
      if (data[i] > maxVal) {
        maxVal = data[i];
        maxIndex = i;
      }
    }

    return { index: maxIndex, confidence: maxVal };
  }

  isReady(): boolean {
    return this.suitModel !== null && this.valueModel !== null;
  }
}

export const cardModelService = new CardModelService();
