import * as tf from '@tensorflow/tfjs';
import { io } from '@tensorflow/tfjs-core';
import LearningInterface from './base';
import  DataManager from '../utils/data_manager';

class DataModel implements LearningInterface {
  epochs: number;
  batchSize: number;
  learningRate: number;
  validateRate: number;
  isTrainedDone: boolean;
  limitSize: number;
  isRunning: boolean;
  isReady: boolean;
  model: tf.LayersModel | null = null;
  dataManager = new DataManager();

  constructor({ 
    epochs = 50,
    batchSize = 16,
    limitSize = 2,
    learningRate = 0.001,
    validateRate = 0.15,
  }: { epochs?: number, batchSize?: number, limitSize?: number, learningRate?: number, validateRate?: number} = {}) {
    this.epochs = epochs;
    this.batchSize = batchSize;
    this.learningRate = learningRate;
    this.validateRate = validateRate;
    this.isRunning = false;
    this.isReady = false;
    this.isTrainedDone = false;
    this.limitSize = limitSize;
  }

  onProgress(progress: number): void {
    console.log(`Training progress: ${progress}%`);
  }

  onLoss(loss: number): void {
    console.log(`Loss: ${loss}`);
  }

  onTrainBegin(log: any): void {
    console.log('Training started', log);
  }

  onTrainEnd(log: any): void {
    console.log('Training ended', log);
  }

  onEpochEnd(epoch: number, logs: any): void {
    console.log(`Epoch ${epoch} ended with logs:`, logs);
  }

  async addData(label: string, data: number[]): Promise<void> {
    const normalizedData = this.normalizeData(data);
    const augmentedData = this.augmentData(normalizedData);
    await this.dataManager.addData(label, augmentedData);
  }

  augmentData(data: number[]): number[] {
    return data.map(value => value + (Math.random() * 0.1 - 0.05)); // 약간의 노이즈 추가
  }

  normalizeData(data: number[]): number[] {
    const mean = data.reduce((sum, value) => sum + value, 0) / data.length;
    const std = Math.sqrt(data.reduce((sum, value) => sum + Math.pow(value - mean, 2), 0) / data.length);
    return data.map(value => (value - mean) / std);
  }

  async train(): Promise<tf.History> {
    if (this.isRunning) {
      throw new Error('Training is already in progress');
    }

    const { xs, ys } = this.dataManager.convertToTensors();

    const jointModel = tf.sequential();
    jointModel.add(tf.layers.dense({
      units: 32,
      activation: 'relu',
      inputShape: [xs.shape[1]],
      kernelRegularizer: tf.regularizers.l2({ l2: 0.01 }) // L2 정규화 추가
    }));
    jointModel.add(tf.layers.dropout({ rate: 0.5 })); // 드롭아웃 레이어 추가
    jointModel.add(tf.layers.dense({
      units: 16,
      activation: 'relu',
      kernelRegularizer: tf.regularizers.l2({ l2: 0.01 }) // L2 정규화 추가
    }));
    jointModel.add(tf.layers.dropout({ rate: 0.5 })); // 드롭아웃 레이어 추가
    jointModel.add(tf.layers.dense({ units: Object.keys(this.dataManager.getLabelMap()).length, activation: 'softmax' }));



    jointModel.compile({
      optimizer: tf.train.adam(this.learningRate), // 학습률 조정
      loss: 'sparseCategoricalCrossentropy',
      metrics: ['accuracy']
    });
    this.model = jointModel;
    this.isRunning = true;

    const history = await jointModel.fit(xs, ys, {
        epochs: this.epochs,
        batchSize: this.batchSize,
        validationSplit: 0.2,
        callbacks: {
          onEpochEnd: (epoch, logs) => this.onEpochEnd(epoch, logs),
          onBatchEnd: (batch: any, logs: any) => {
          },
          onTrainBegin: (logs) => this.onTrainBegin(logs),
          onTrainEnd: (logs) => {
            this.isRunning = false;
            this.isReady = true;
            this.onTrainEnd(logs);
          },
        }
    });
    return history;
  }

  public async infer(data: number[]): Promise<Map<string, number>> {
    if (!this.isReady || !this.model) {
      throw new Error('Model is not ready');
    }

    // 입력 데이터를 텐서로 변환
    const input = tf.tensor2d([data], [1, data.length], 'float32');

    try {
      // 모델 예측 수행
      const prediction = this.model.predict(input) as tf.Tensor<tf.Rank>;

      // 예측 결과 값들을 추출
      const values = await prediction.data();

      // 클래스 이름과 확률을 저장할 Map 생성
      const classProbabilities = new Map<string, number>();
      const EPSILON = 1e-6; // 매우 작은 값을 표현하기 위한 엡실론

      // 라벨 맵을 가져옴
      const labelMap = this.dataManager.getLabelMap();
      const reverseLabelMap = Object.entries(labelMap).reduce((acc, [key, value]) => {
        acc[value] = key;
        return acc;
      }, {} as { [key: number]: string });

      // 각 클래스의 확률 계산 및 조정
      for (let i = 0; i < values.length; i++) {
        let probability = Math.max(0, Math.min(1, values[i])); // 확률 값을 0과 1 사이로 조정
        probability = probability < EPSILON ? 0 : probability; // 매우 작은 확률 값을 0으로 간주
        const className = reverseLabelMap[i]; // 클래스 이름을 문자열로 가져옴
        if (className) {
          const existingProbability = classProbabilities.get(className);
          if (existingProbability !== undefined) {
            classProbabilities.set(className, existingProbability + probability);
          } else {
            classProbabilities.set(className, probability);
          }
        }
      }

      // 텐서 해제
      prediction.dispose();
      input.dispose();

      console.log('classProbabilities', classProbabilities);
      return classProbabilities;
    } catch (error) {
      // 텐서 해제
      input.dispose();
      throw error;
    }
  }
  
  
  
  

  async saveModel(handlerOrURL: io.IOHandler | string, config?: io.SaveConfig): Promise<void> {
    if (!this.model) {
      throw new Error('No model to save');
    }

    await this.model.save(handlerOrURL, config);
  }

  running(): boolean {
    return this.isRunning;
  }

  ready(): boolean {
    return this.isReady;
  }

  async load({ jsonURL, labels }: { jsonURL: string, labels: Array<string> }): Promise<void> {
    this.model = await tf.loadLayersModel(jsonURL);
    this.isReady = true;

    // Assuming the labels are provided in the correct order
    this.dataManager = new DataManager();
    labels.forEach((label, index) => {
      this.dataManager.getLabelMap()[label] = index;
    });
  }
}

export default DataModel;