import * as tf from "@tensorflow/tfjs-node";
import { readFile } from "node:fs/promises";
import { existsSync } from "node:fs";
import { join } from "node:path";

async function train() {
  const savePath = process.argv[2] || 'sample/model';
  const data = JSON.parse(await readFile(join(process.argv[3] || 'sample/dataset', 'dataset.json'), "utf-8"));
  const X = tf.tensor2d(data.map((d: any) => d.x.map((v: any, i: number) => v / (i == 0 ? 20 : i == 1 ? 200 : 1))));
  const y = tf.tensor2d(data.map((d: any) => [d.y]));

  let model: tf.LayersModel;
  const modelJsonPath = join(savePath, "model.json");

  if (existsSync(modelJsonPath)) {
    console.log("기존 모델을 불러와 추가 학습을 진행합니다.");
    model = await tf.loadLayersModel(`file://${modelJsonPath}`);
  } else {
    console.log("새 모델을 생성합니다.");
    model = tf.sequential({
      layers: [
        tf.layers.dense({ units: 32, activation: 'relu', inputShape: [5] }),
        tf.layers.dense({ units: 16, activation: 'relu' }),
        tf.layers.dense({ units: 1 })
      ]
    });
  }

  model.compile({ optimizer: 'adam', loss: 'meanSquaredError' });
  const epochs = parseInt(process.argv[4]) || 50;
  await model.fit(X, y, { epochs: epochs, verbose: 0 });
  await model.save(`file://${savePath}`);
  console.log(`학습 완료: 모델이 ${savePath}에 저장되었습니다.`);
}
train().catch(console.error);