package com.kamco.cd.training.postgres.repository.model;

import static com.kamco.cd.training.postgres.entity.QDatasetEntity.datasetEntity;
import static com.kamco.cd.training.postgres.entity.QDatasetObjEntity.datasetObjEntity;
import static com.kamco.cd.training.postgres.entity.QModelDatasetEntity.modelDatasetEntity;
import static com.kamco.cd.training.postgres.entity.QModelDatasetMappEntity.modelDatasetMappEntity;
import static com.kamco.cd.training.postgres.entity.QModelHyperParamEntity.modelHyperParamEntity;
import static com.kamco.cd.training.postgres.entity.QModelMasterEntity.modelMasterEntity;
import static com.kamco.cd.training.postgres.entity.QModelMetricsTestEntity.modelMetricsTestEntity;
import static com.kamco.cd.training.postgres.entity.QModelMetricsTrainEntity.modelMetricsTrainEntity;
import static com.kamco.cd.training.postgres.entity.QModelMetricsValidationEntity.modelMetricsValidationEntity;

import com.kamco.cd.training.common.enums.DetectionClassification;
import com.kamco.cd.training.common.enums.ModelType;
import com.kamco.cd.training.common.enums.TrainStatusType;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.DetailSummary;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.HyperSummary;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.MappingDataset;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.ModelBestEpoch;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.ModelFileInfo;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.ModelTestMetrics;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.ModelTrainMetrics;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.ModelValidationMetrics;
import com.kamco.cd.training.model.dto.ModelTrainDetailDto.TransferHyperSummary;
import com.kamco.cd.training.model.dto.ModelTrainMngDto.ModelProgressStepDto;
import com.kamco.cd.training.postgres.entity.ModelMasterEntity;
import com.kamco.cd.training.postgres.entity.QModelHyperParamEntity;
import com.kamco.cd.training.postgres.entity.QModelMasterEntity;
import com.querydsl.core.types.Expression;
import com.querydsl.core.types.Projections;
import com.querydsl.core.types.dsl.CaseBuilder;
import com.querydsl.jpa.JPAExpressions;
import com.querydsl.jpa.impl.JPAQueryFactory;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Repository;

@Slf4j
@Repository
@RequiredArgsConstructor
public class ModelDetailRepositoryImpl implements ModelDetailRepositoryCustom {

  private final JPAQueryFactory queryFactory;

  /**
   * 모델 조회
   *
   * @param uuid
   * @return
   */
  @Override
  public Optional<ModelMasterEntity> findByUuid(UUID uuid) {
    return Optional.ofNullable(
        queryFactory
            .select(modelMasterEntity)
            .from(modelMasterEntity)
            .where(modelMasterEntity.uuid.eq(uuid))
            .fetchOne());
  }

  @Override
  public DetailSummary getModelDetailSummary(UUID uuid) {
    QModelMasterEntity beforeModel = new QModelMasterEntity("beforeModel"); // alias

    Expression<UUID> beforeModelUuid =
        com.querydsl.jpa.JPAExpressions.select(beforeModel.uuid)
            .from(beforeModel)
            .where(beforeModel.id.eq(modelMasterEntity.beforeModelId));

    return queryFactory
        .select(
            Projections.constructor(
                DetailSummary.class,
                modelMasterEntity.id,
                modelMasterEntity.uuid,
                modelMasterEntity.modelNo,
                modelMasterEntity.modelVer,
                modelMasterEntity.step1StrtDttm,
                modelMasterEntity.step1EndDttm,
                modelMasterEntity.step2StrtDttm,
                modelMasterEntity.step2EndDttm,
                modelMasterEntity.statusCd,
                modelMasterEntity.trainType,
                beforeModelUuid))
        .from(modelMasterEntity)
        .where(modelMasterEntity.uuid.eq(uuid))
        .fetchOne();
  }

  @Override
  public HyperSummary getByModelHyperParamSummary(UUID uuid) {
    return queryFactory
        .select(
            Projections.constructor(
                HyperSummary.class,
                modelHyperParamEntity.uuid,
                modelHyperParamEntity.id,
                modelHyperParamEntity.hyperVer,
                modelHyperParamEntity.backbone,
                modelHyperParamEntity.inputSize,
                modelHyperParamEntity.cropSize,
                modelHyperParamEntity.batchSize))
        .from(modelHyperParamEntity)
        .where(
            modelHyperParamEntity.id.eq(
                JPAExpressions.select(modelMasterEntity.hyperParamId)
                    .from(modelMasterEntity)
                    .where(modelMasterEntity.uuid.eq(uuid))))
        .fetchOne();
  }

  @Override
  public TransferHyperSummary getByModelTransferHyperParamSummary(UUID uuid) {

    QModelMasterEntity subMaster = new QModelMasterEntity("subMaster");
    QModelHyperParamEntity subHyper = new QModelHyperParamEntity("subHyper");

    return queryFactory
        .select(
            Projections.constructor(
                TransferHyperSummary.class,
                modelHyperParamEntity.uuid,
                modelHyperParamEntity.id,
                modelHyperParamEntity.hyperVer,
                modelHyperParamEntity.backbone,
                modelHyperParamEntity.inputSize,
                modelHyperParamEntity.cropSize,
                modelHyperParamEntity.batchSize,
                subHyper.uuid,
                subHyper.id,
                subHyper.hyperVer,
                subHyper.backbone,
                subHyper.inputSize,
                subHyper.cropSize,
                subHyper.batchSize))
        .from(modelMasterEntity)
        .innerJoin(modelHyperParamEntity)
        .on(modelHyperParamEntity.id.eq(modelMasterEntity.hyperParamId))
        .leftJoin(subMaster)
        .on(subMaster.id.eq(modelMasterEntity.beforeModelId))
        .leftJoin(subHyper)
        .on(subHyper.id.eq(subMaster.hyperParamId))
        .where(modelMasterEntity.uuid.eq(uuid))
        .fetchOne();
  }

  @Override
  public List<MappingDataset> getByModelMappingDataset(UUID uuid) {
    return queryFactory
        .select(
            Projections.constructor(
                MappingDataset.class,
                modelMasterEntity.id,
                datasetEntity.id,
                datasetEntity.dataType,
                datasetEntity.compareYyyy,
                datasetEntity.targetYyyy,
                datasetEntity.roundNo,

                // G1 - building
                new CaseBuilder()
                    .when(
                        modelMasterEntity
                            .modelNo
                            .eq(ModelType.G1.getId())
                            .and(
                                datasetObjEntity.targetClassCd.eq(
                                    DetectionClassification.BUILDING.getId())))
                    .then(1L)
                    .otherwise(0L)
                    .sum(),

                // G1 - container
                new CaseBuilder()
                    .when(
                        modelMasterEntity
                            .modelNo
                            .eq(ModelType.G1.getId())
                            .and(
                                datasetObjEntity.targetClassCd.eq(
                                    DetectionClassification.CONTAINER.getId())))
                    .then(1L)
                    .otherwise(0L)
                    .sum(),

                // G2 - waste
                new CaseBuilder()
                    .when(
                        modelMasterEntity
                            .modelNo
                            .eq(ModelType.G2.getId())
                            .and(
                                datasetObjEntity.targetClassCd.eq(
                                    DetectionClassification.WASTE.getId())))
                    .then(1L)
                    .otherwise(0L)
                    .sum(),

                // G3 - 나머지
                new CaseBuilder()
                    .when(
                        modelMasterEntity
                            .modelNo
                            .eq(ModelType.G3.getId())
                            .and(
                                datasetObjEntity
                                    .targetClassCd
                                    .isNotNull()
                                    .and(
                                        datasetObjEntity.targetClassCd.notIn(
                                            DetectionClassification.BUILDING.getId(),
                                            DetectionClassification.CONTAINER.getId(),
                                            DetectionClassification.WASTE.getId(),
                                            DetectionClassification.SOLAR.getId()))))
                    .then(1L)
                    .otherwise(0L)
                    .sum(),

                // G4 - solar
                new CaseBuilder()
                    .when(
                        modelMasterEntity
                            .modelNo
                            .eq(ModelType.G4.getId())
                            .and(
                                datasetObjEntity.targetClassCd.eq(
                                    DetectionClassification.SOLAR.getId())))
                    .then(1L)
                    .otherwise(0L)
                    .sum()))
        .from(modelMasterEntity)
        .innerJoin(modelDatasetEntity)
        .on(modelMasterEntity.id.eq(modelDatasetEntity.model.id))
        .innerJoin(modelDatasetMappEntity)
        .on(modelMasterEntity.id.eq(modelDatasetMappEntity.modelUid))
        .innerJoin(datasetEntity)
        .on(modelDatasetMappEntity.datasetUid.eq(datasetEntity.id))
        .leftJoin(datasetObjEntity)
        .on(datasetEntity.id.eq(datasetObjEntity.datasetUid))
        .where(modelMasterEntity.uuid.eq(uuid))
        .groupBy(
            modelMasterEntity.id,
            datasetEntity.id,
            datasetEntity.dataType,
            datasetEntity.compareYyyy,
            datasetEntity.targetYyyy,
            datasetEntity.roundNo)
        .fetch();
  }

  @Override
  public ModelMasterEntity findByModelByUUID(UUID uuid) {
    return queryFactory
        .selectFrom(modelMasterEntity)
        .where(modelMasterEntity.uuid.eq(uuid))
        .fetchOne();
  }

  @Override
  public List<ModelTrainMetrics> getModelTrainMetricResult(UUID uuid) {
    ModelMasterEntity modelMasterEntity = findByModelByUUID(uuid);
    if (modelMasterEntity == null) {
      return List.of();
    }

    return queryFactory
        .select(
            Projections.constructor(
                ModelTrainMetrics.class,
                modelMetricsTrainEntity.epoch,
                modelMetricsTrainEntity.iteration,
                modelMetricsTrainEntity.loss,
                modelMetricsTrainEntity.lr,
                modelMetricsTrainEntity.durationTime))
        .from(modelMetricsTrainEntity)
        .where(modelMetricsTrainEntity.model.id.eq(modelMasterEntity.getId()))
        .fetch();
  }

  @Override
  public List<ModelValidationMetrics> getModelValidationMetricResult(UUID uuid) {
    ModelMasterEntity modelMasterEntity = findByModelByUUID(uuid);
    if (modelMasterEntity == null) {
      return List.of();
    }

    return queryFactory
        .select(
            Projections.constructor(
                ModelValidationMetrics.class,
                modelMetricsValidationEntity.epoch,
                modelMetricsValidationEntity.aAcc,
                modelMetricsValidationEntity.mFscore,
                modelMetricsValidationEntity.mPrecision,
                modelMetricsValidationEntity.mRecall,
                modelMetricsValidationEntity.mIou,
                modelMetricsValidationEntity.mAcc,
                modelMetricsValidationEntity.changedFscore,
                modelMetricsValidationEntity.changedPrecision,
                modelMetricsValidationEntity.changedRecall,
                modelMetricsValidationEntity.unchangedFscore,
                modelMetricsValidationEntity.unchangedPrecision,
                modelMetricsValidationEntity.unchangedRecall))
        .from(modelMetricsValidationEntity)
        .where(modelMetricsValidationEntity.model.id.eq(modelMasterEntity.getId()))
        .fetch();
  }

  @Override
  public List<ModelTestMetrics> getModelTestMetricResult(UUID uuid) {
    ModelMasterEntity modelMasterEntity = findByModelByUUID(uuid);
    if (modelMasterEntity == null) {
      return List.of();
    }

    return queryFactory
        .select(
            Projections.constructor(
                ModelTestMetrics.class,
                modelMetricsTestEntity.model1,
                modelMetricsTestEntity.tp,
                modelMetricsTestEntity.fp,
                modelMetricsTestEntity.fn,
                modelMetricsTestEntity.precisions,
                modelMetricsTestEntity.recall,
                modelMetricsTestEntity.f1Score,
                modelMetricsTestEntity.accuracy,
                modelMetricsTestEntity.iou,
                modelMetricsTestEntity.detectionCount,
                modelMetricsTestEntity.gtCount))
        .from(modelMetricsTestEntity)
        .where(modelMetricsTestEntity.model.id.eq(modelMasterEntity.getId()))
        .fetch();
  }

  @Override
  public ModelBestEpoch getModelTrainBestEpoch(UUID uuid) {
    ModelMasterEntity modelMasterEntity = findByModelByUUID(uuid);
    if (modelMasterEntity == null) {
      return null;
    }

    return queryFactory
        .select(
            Projections.constructor(
                ModelBestEpoch.class,
                modelMetricsTrainEntity.epoch,
                modelMetricsTrainEntity.loss,
                modelMetricsValidationEntity.mFscore,
                modelMetricsValidationEntity.mPrecision,
                modelMetricsValidationEntity.mRecall,
                modelMetricsValidationEntity.mIou,
                modelMetricsValidationEntity.mAcc))
        .from(modelMetricsTrainEntity)
        .leftJoin(modelMetricsValidationEntity)
        .on(
            modelMetricsTrainEntity.model.eq(modelMetricsValidationEntity.model),
            modelMetricsTrainEntity.epoch.eq(modelMetricsValidationEntity.epoch))
        .where(
            modelMetricsTrainEntity.model.id.eq(modelMasterEntity.getId()),
            modelMetricsTrainEntity.epoch.eq(modelMasterEntity.getBestEpoch()))
        .fetchOne();
  }

  @Override
  public ModelFileInfo getModelTrainFileInfo(UUID uuid) {
    return queryFactory
        .select(
            Projections.constructor(
                ModelFileInfo.class,
                modelMasterEntity
                    .packingState
                    .eq(TrainStatusType.COMPLETED.getId())
                    .coalesce(false),
                modelMasterEntity.modelVer))
        .from(modelMasterEntity)
        .where(modelMasterEntity.uuid.eq(uuid))
        .fetchOne();
  }

  @Override
  public List<ModelProgressStepDto> findModelTrainProgressInfo(UUID uuid) {
    ModelMasterEntity entity = findByModelByUUID(uuid);
    if (entity == null) {
      return List.of();
    }

    List<ModelProgressStepDto> steps = new ArrayList<>();

    // 0단계 : 대기 상태
    steps.add(
        ModelProgressStepDto.builder()
            .step(0)
            .status(TrainStatusType.READY.getId())
            .startTime(entity.getCreatedDttm())
            .endTime(null)
            .isError(false)
            .build());

    // 1단계 : Train/Validation 실행
    boolean step1Active =
        entity.getStep1StrtDttm() != null
            && !TrainStatusType.READY.getId().equals(entity.getStep1State());

    if (step1Active) {
      steps.add(
          ModelProgressStepDto.builder()
              .step(1)
              .status(entity.getStep1State())
              .startTime(entity.getStep1StrtDttm())
              .endTime(entity.getStep1EndDttm())
              .isError(TrainStatusType.ERROR.getId().equals(entity.getStep1State()))
              .build());
    }

    // 2단계 : Test 실행
    boolean step2Done = entity.getStep2State() != null;

    if (step2Done) {
      steps.add(
          ModelProgressStepDto.builder()
              .step(2)
              .status(entity.getStep2State())
              .startTime(entity.getStep2StrtDttm())
              .endTime(entity.getStep2EndDttm())
              .isError(TrainStatusType.ERROR.getId().equals(entity.getStep2State()))
              .build());
    }

    // 3단계 : 패키징
    boolean step3Done = entity.getPackingState() != null;

    if (step3Done) {
      steps.add(
          ModelProgressStepDto.builder()
              .step(3)
              .status(entity.getPackingState())
              .startTime(entity.getPackingStrtDttm())
              .endTime(entity.getPackingEndDttm())
              .isError(TrainStatusType.ERROR.getId().equals(entity.getPackingState()))
              .build());
    }

    return steps;
  }

  @Override
  public ModelMasterEntity findByModelBeforeId(Long beforeModelId) {
    return queryFactory
        .selectFrom(modelMasterEntity)
        .where(modelMasterEntity.id.eq(beforeModelId))
        .fetchOne();
  }
}
