///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
// mobilenet 모델을 이용한 전이학습 방법 
///////////////////////////////////////////////////////////////////////////

import * as tf from '@tensorflow/tfjs';
import { setWasmPaths } from '@tensorflow/tfjs-backend-wasm';
import { dispose } from '@tensorflow/tfjs';
import { io } from '@tensorflow/tfjs-core';
import LearningInterface from './base';
import { capture, loadModel, isTensor } from '../utils/tf';
import { cropTo } from '../utils/canvas';
import { flatOneHot, fisherYates, Sample } from '../utils/dataset';
import { Initializer } from '@tensorflow/tfjs-layers/dist/initializers';


class LearningMobilenet implements LearningInterface {
  model: tf.LayersModel | null;
  epochs: number;
  batchSize: number;
  learningRate: number;
  validateRate: number;
  isRunning: boolean;
  isReady: boolean;
  isTrainedDone: boolean;
  limitSize: number;
  mobilenetModule: tf.LayersModel | null;
  imageExamples: Float32Array[][] = [];
  classNumber: string[] = [];
  
  readonly MOBILE_NET_INPUT_WIDTH = 224;
  readonly MOBILE_NET_INPUT_HEIGHT = 224;
  readonly MOBILE_NET_INPUT_CHANNEL = 3;
  readonly IMAGE_NORMALIZATION_FACTOR = 255.0;
  

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

  }

  // 진행 상태를 나타내는 이벤트를 정의합니다.
  public onProgress: (progress: number) => void = () => {};

  public onLoss: (loss: number) => void = () => {};

  public onEvents: (logs: any) => void = () => {};

  public onTrainBegin: (log: any) => void = () => {};

  public onTrainEnd: (log: any) => void = () => {};

  public onEpochEnd: (epoch: number, logs: any) => void = () => {};
  
  //
  // 기존의 모델 로드
  public async load({jsonURL, labels}: {jsonURL :string, labels : Array<string>}): Promise<void> {
    if (labels.length <= 0) {
      return Promise.reject(new Error('Labels length is 0'));
    }

    try {
      this.model = await tf.loadLayersModel(jsonURL);
      for(var i = 0; i < labels.length; i++) {
        this.registerClassNumber(labels[i]);
    }
      this.isReady = true;
      this.model.summary();
    } catch (error) {
      console.error('Model load failed', error);
      throw error;
    }
  }


  private registerClassNumber(value: string): number {
    // 중복 값의 인덱스를 찾습니다.
    const existingIndex = this.classNumber.indexOf(value);

    // 중복 값이 있다면 해당 인덱스를 반환합니다.
    if (existingIndex !== -1) {
        return existingIndex;
    }

    // 중복 값이 없다면 새로운 항목을 추가하고 그 인덱스를 반환합니다.
    this.classNumber.push(value);
    return this.classNumber.length - 1;
  }

  private _convertToTfDataset() {
    for (let i = 0; i < this.imageExamples.length; i++) {
        this.imageExamples[i] = fisherYates(this.imageExamples[i]) as Float32Array[];
    }

    const trainDataset: Sample[] = [];
    const validationDataset: Sample[] = [];
    
    for (let i = 0; i < this.imageExamples.length; i++) {
        const classLength = this.imageExamples[i].length;

        // 클래스의 전체 데이터 수를 사용하여 학습 및 검증 데이터 수 계산
        const numValidation = Math.ceil(this.validateRate * classLength);
        const numTrain = classLength - numValidation;

        // One-Hot 인코딩을 사용하여 라벨 생성
        const y = flatOneHot(i, this.classNumber.length);

        // numTrain과 numValidation에 따라 데이터를 학습 및 검증 데이터로 분할
        const classTrain: Sample[] = this.imageExamples[i].slice(0, numTrain).map(dataArray => ({ data: dataArray, label: y }));
        trainDataset.push(...classTrain);
        
        const classValidation: Sample[] = this.imageExamples[i].slice(numTrain, numTrain + numValidation).map(dataArray => ({ data: dataArray, label: y }));
        validationDataset.push(...classValidation);
    }
    
    // Shuffle entire datasets
    const shuffledTrainDataset = fisherYates(trainDataset) as Sample[];
    const shuffledValidationDataset = fisherYates(validationDataset) as Sample[];

    // Convert to tf.data.Dataset
    const trainX = tf.data.array(shuffledTrainDataset.map(sample => sample.data));
    const validationX = tf.data.array(shuffledValidationDataset.map(sample => sample.data));
    const trainY = tf.data.array(shuffledTrainDataset.map(sample => sample.label));
    const validationY = tf.data.array(shuffledValidationDataset.map(sample => sample.label));

    return {
        trainDataset: tf.data.zip({ xs: trainX, ys: trainY }),
        validationDataset: tf.data.zip({ xs: validationX, ys: validationY })
    };
}


  // 학습 데이타 등록 
  public async addData(label: string, data: any): Promise<void> {
    try {
      if (this.mobilenetModule !== null) {
        const cap = isTensor(data) ? data : await capture(data, false);

        const predict = tf.tidy(() => {
          return this.mobilenetModule!.predict(cap) as tf.Tensor;
        });

        const activation = await predict.data() as Float32Array;

        const classIndex = this.registerClassNumber(label);
        if (!this.imageExamples[classIndex]) {
          this.imageExamples[classIndex] = [];
        }
        this.imageExamples[classIndex].push(activation);

        if(this.classNumber.length >= this.limitSize) {
          this.isReady = true;
        }

        // Dispose of the prediction tensor to free memory
        predict.dispose();

        // If cap is not a tensor, we don't need to dispose it. Otherwise, we should.
        if (!isTensor(data)) {
          cap.dispose();
        }

      } else {
        throw new Error('mobilenetModule is null');
      }
      return Promise.resolve();
    } catch (error) {
      console.error('Model training failed', error);
      throw error;
    }
  }

  public async init() {
    try {
      console.log('init call')
      await this.setupBackend();
      this.mobilenetModule = await loadModel();
    } catch(error) {
      console.log('init Error', error);
      throw error;
    } 
  } 

  private async setupBackend() {
    const isWasmSupported = await this.checkWasmSupport();
    if (isWasmSupported) {
      console.log('Backend is set to WebAssembly');
    } else {
      await tf.setBackend('cpu');
      await tf.ready();
      console.log('Backend is set to CPU');
    }
  }

  private async checkWasmSupport(): Promise<boolean> {
    try {
      const wasmVersion = (tf.version as Record<string, string>)['tfjs-backend-wasm'] || tf.version['tfjs-core'];
      setWasmPaths(`https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-wasm@${wasmVersion}/dist/`);
      await tf.setBackend('wasm');
      await tf.ready();
      return true;
    } catch (error) {
      console.warn('WASM backend is not supported in this environment.', error);
      return false;
    }
  }

  // 모델 학습 처리
  public async train(): Promise<tf.History> {
    if (this.isRunning) {
      return Promise.reject(new Error('Training is already in progress.'));
    }
    
    // 콜백 정의
    const customCallback = {
      onTrainBegin: (log: any) => {
        this.isTrainedDone = false;
        this.onTrainBegin(log);
      },
      onTrainEnd: (log: any) => {
        this.isTrainedDone = true;
        this.onTrainEnd(log);
        this.isRunning = false;
      },
      onBatchBegin: (batch: any, logs: any) => {
        //console.log(`Batch ${batch} is starting.`);
      },
      onBatchEnd: (batch: any, logs: any) => {
        //console.log(`Batch ${batch} has ended.`);
      },
      onEpochBegin: (epoch: number, logs: any) => {
        //console.log(`Epoch ${epoch+1} is starting.`, logs);
      },
      onEpochEnd: (epoch: number, logs: any) => {
        this.onEpochEnd(epoch, logs);
        this.onLoss(logs.loss);
        this.onProgress(epoch+1);
        this.onEvents({
          epoch: epoch,
          logs: logs,
        });
      }
    };

    try {
      this.isRunning = true;
      if (this.classNumber.length < this.limitSize) {
        return Promise.reject(new Error('Please train Data need over 2 data length'));
      }
      const datasets = this._convertToTfDataset();
      const trainData = datasets.trainDataset.batch(this.batchSize);
      const validationData = datasets.validationDataset.batch(this.batchSize);
      const optimizer = tf.train.adam(this.learningRate); 

      const trainModel = await this._createModel(optimizer);

      const jointModel = tf.sequential();
      jointModel.add(this.mobilenetModule!);
      jointModel.add(trainModel);
      this.model = jointModel;

      const history = await trainModel.fitDataset(trainData, {
        epochs: this.epochs,
        validationData: validationData,
        callbacks: customCallback
      });

      optimizer.dispose();
      return history;
    } catch (error) {
      this.isRunning = false;
      console.error('Model training failed', error);
      throw error;
    }
  }
  
  // 추론하기
  public async infer(data: any): Promise<Map<string, number>> {
    if (this.model === null) {
      return Promise.reject(new Error('Model is Null'));
    }
  
    try {
      const classProbabilities = new Map<string, number>();
      const croppedImage = cropTo(data, 224, false);
      const captured = await capture(croppedImage, false);

      const logits = tf.tidy(() => {
        return this.model!.predict(captured);
      });

      const values = await (logits as tf.Tensor<tf.Rank>).data();
      const EPSILON = 1e-6; // 매우 작은 값을 표현하기 위한 엡실론

      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 = this.classNumber[i]; // 클래스 이름
        const existingProbability = classProbabilities.get(className);
        if (existingProbability !== undefined) {
          classProbabilities.set(className, existingProbability + probability);
        } else {
          classProbabilities.set(className, probability);
        }
      }

      console.log('classProbabilities', classProbabilities);
      dispose(logits);
      return classProbabilities;
    } catch (error) {
      throw error;
    }
  }

  
  // 모델 저장
  public async saveModel(handlerOrURL: io.IOHandler | string, config?: io.SaveConfig): Promise<void> {
    try {
      console.log('saveModel try', this.isTrainedDone);
      if (!this.isTrainedDone) {
        return Promise.reject(new Error('Train is not done status'));
      }
      await this.model?.save(handlerOrURL, config);
    } catch(e) {
      console.log('saveModel Error', e);
    }
  }

  // 진행중 여부
  public running(): boolean {
    return this.isRunning;
  }


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

  // 모델 저장
  private async _createModel(optimizer: tf.AdamOptimizer): Promise<tf.Sequential> {
    try {
      // 입력 이미지 크기에 맞게 모델 구조 수정
      let varianceScaling: Initializer;
      varianceScaling = tf.initializers.varianceScaling({});

      const trainModel = tf.sequential({
        layers: [
          tf.layers.dense({
            inputShape: this.mobilenetModule!.outputs[0].shape.slice(1),
            units: 128,
            activation: 'relu',
            kernelInitializer: varianceScaling, // 'varianceScaling'
            useBias: true
          }),
          tf.layers.dense({
            kernelInitializer: varianceScaling, // 'varianceScaling'
            useBias: false,
            activation: 'softmax',
            units: this.classNumber.length, 
          })
        ]
      });
      
      trainModel.compile({
        loss: 'categoricalCrossentropy',
        optimizer: optimizer,
        metrics: ['accuracy']
      });
      
      return trainModel;
    } catch (error) {
      console.error('Failed to load model', error);
      throw error;
    }
  }
}

export default LearningMobilenet;