import {
  ChangeDetectionStrategy,
  Component,
  inject,
  Input,
  OnChanges,
  OnInit,
  SimpleChanges,
} from '@angular/core';
import {
  type GeneralAnalysis,
  type InputPipelineAnalysis,
  type SimpleDataTable,
} from 'org_xprof/frontend/app/common/interfaces/data_table';
import {
  GeneralProps,
  SummaryInfo,
  SummaryInfoConfig,
} from 'org_xprof/frontend/app/common/interfaces/summary_info';
import {DATA_SERVICE_INTERFACE_TOKEN} from 'org_xprof/frontend/app/services/data_service_v2/data_service_v2_interface';

/** Options for parsing summary info config. */
interface ParseConfigOptions {
  level?: number;
  customInput?: google.visualization.DataObjectCell[];
  baselineProps?: GeneralProps;
  baselineCustomInput?: google.visualization.DataObjectCell[];
  parentGoodMetric?: boolean;
}

/** Parses a string or number value to a number, returning null if invalid. */
export function parseToNumber(
  val: string | number | null | undefined,
): number | null {
  if (val === undefined || val === null || val === '') return null;
  const cleaned = String(val).replace('%', '').trim();
  const num = Number(cleaned);
  return Number.isNaN(num) ? null : num;
}

/**
 * Configuration Assumptions:
 * 1. by default, values/sdv values/nested children values all reads from
 * a k-v pair property with field key specified in the config object.
 * 2. It's required that each record in the same config object (eg.
 * NON_TPU_SUMMARY_INFO) should consume the same property (eg.
 * inputPipelineAnalysisProperty).
 * 3. The summary info data is in nested structure, currently consumes 2 levels
 * of metrics (Top layer level 1 + its children layer)
 */

/**
 * Configuration Instructions
 * 1. In this file we chain multiple config objects to display metrics in order:
 * GENERIC_SUMMARY_INFO_BEFORE + TPU_SUMMARY_INFO/NON_TPU_SUMMARY_INFO +
 * GENERIC_SUMMARY_INFO_AFTER
 * 2. Values read from the k-v pair property could be undefined for oss
 * (dependent on backend version). Be sure to add a fallback logic (eg. || '')
 * when reading it in the parseDataFromConfig function call.
 * 3. Values are by default read from k-v pairs, we can also use a custom
 * callback function (eg. getChildValues) to read data into the propertyValues
 * as a string list.
 * 4. If `customInput` is specified when calling `parseDataFromConfig`, it will
 * be used as input to callback function `getValue` or `getChildValues`.
 */

/** Generic summary info, display on top of the list */
const GENERIC_SUMMARY_INFO_BEFORE: SummaryInfoConfig[] = [
  {
    title: 'Average Step Time',
    goodMetric: false,
    valueKey: 'steptime_ms_average',
    sdvKey: 'steptime_ms_standard_deviation',
    unit: 'ms',
    valueColor: 'red',
    trainingOnly: true,
    childrenInfoConfig: [
      {title: 'Idle', valueKey: 'idle_ms_average', unit: 'ms'},
      {title: 'Input', valueKey: 'input_ms_average', unit: 'ms'},
      {title: 'Compute', valueKey: 'compute_ms_average', unit: 'ms'},
    ],
  },
];

/** Generic summary info, display on bottom of the list */
const GENERIC_SUMMARY_INFO_AFTER: SummaryInfoConfig[] = [
  {
    title: 'Framework Op Placement',
    tooltip: 'Ratio of Framework Ops executed on the host and device.',
    description: 'generally desired to have more ops on device',
    childrenInfoConfig: [
      {title: 'Host', valueKey: 'host_tf_op_percent'},
      {title: 'Device', valueKey: 'device_tf_op_percent'},
    ],
  },
  {
    title: 'Op Time Spent on Eager Execution',
    tooltip:
      'Out of the total op execution time on host (device), excluding idle time, the percentage of which used eager execution.',
    goodMetric: false,
    childrenInfoConfig: [
      {title: 'Host', valueKey: 'host_op_time_eager_percent'},
      {title: 'Device', valueKey: 'device_op_time_eager_percent'},
    ],
  },
  {
    title: 'Device Compute Precisions',
    description: 'out of Total Device Time',
    childrenInfoConfig: [
      {title: '16-bit', valueKey: 'device_compute_16bit_percent'},
      {title: '32-bit', valueKey: 'device_compute_32bit_percent'},
    ],
  },
];

const NON_TPU_SUMMARY_INFO: SummaryInfoConfig[] = [
  {
    title: 'All Others Time',
    unit: 'ms',
    valueKey: 'other_time_ms_avg',
    sdvKey: 'other_time_ms_sdv',
  },
  {
    title: 'Compilation Time',
    unit: 'ms',
    valueKey: 'compile_time_ms_avg',
    sdvKey: 'compile_time_ms_sdv',
  },
  {
    title: 'Output Time',
    unit: 'ms',
    valueKey: 'outfeed_time_ms_avg',
    sdvKey: 'outfeed_time_ms_sdv',
  },
  {
    title: 'Input Time',
    unit: 'ms',
    valueKey: 'infeed_time_ms_avg',
    sdvKey: 'infeed_time_ms_sdv',
  },
  {
    title: 'Kernel Launch Time',
    unit: 'ms',
    valueKey: 'kernel_launch_time_ms_avg',
    sdvKey: 'kernel_launch_time_ms_sdv',
  },
  {
    title: 'Host Compute Time',
    unit: 'ms',
    valueKey: 'host_compute_time_ms_avg',
    sdvKey: 'host_compute_time_ms_sdv',
  },
  {
    title: 'Device Collective Communication Time',
    unit: 'ms',
    valueKey: 'device_collectives_time_ms_avg',
    sdvKey: 'device_collectives_time_ms_sdv',
  },
  {
    title: 'Device to Device Time',
    unit: 'ms',
    valueKey: 'device_to_device_time_ms_avg',
    sdvKey: 'device_to_device_time_ms_sdv',
  },
  {
    title: 'Device Compute Time',
    unit: 'ms',
    valueKey: 'device_compute_time_ms_avg',
    sdvKey: 'device_compute_time_ms_sdv',
  },
];

const TPU_SUMMARY_INFO: SummaryInfoConfig[] = [
  {
    title: 'FLOPS Utilization',
    goodMetric: true,
    tooltip:
      "Why two numbers: The first number shows the hardware utilization based on the hardware performance counter. The second one shows the performance compared to the program's optimal performance considering the instruction mix (i.e., the ratio of floating-point operations and memory operations).",
    childrenInfoConfig: [
      {
        title: 'Utilization of TPU Matrix Units',
        valueKey: 'mxu_utilization_percent',
      },
      {
        title: "Compared to Program's Optimal FLOPS",
        valueKey: 'flop_rate_utilization_relative_to_roofline',
      },
    ],
  },
  {
    title: 'TPU Duty Cycle',
    tooltip: 'Percentage of the device time that is busy.',
    goodMetric: true,
    valueKey: 'device_duty_cycle_percent',
  },
  {
    title: 'Memory Bandwidth Utilization',
    tooltip:
      "Why two numbers: The first number shows the memory bandwidth utilization based on the hardware performance counters for HBM. The second one shows the bandwidth compared to the program's optimal performance considering the bytes accessed based on the input/output size and data type.",
    goodMetric: true,
    childrenInfoConfig: [
      {
        title: 'HBM Utilization',
        valueKey: 'hbm_utilization_percent',
      },
      {
        title: "Compared to Program's Optimal Bandwidth (VMem + SpMem + HBM)",
        valueKey: 'memory_bw_utilization_relative_to_hw_limit',
      },
    ],
  },
  {
    title: 'Firmware Power Metrics (Power/Timescale)',
    tooltip:
      'Max power consumption of different components/rails with different timescale.',
    getChildValues: (props) =>
      (props as GeneralProps)['fw_power_metrics']?.split('##') || [],
  },
  {
    title: 'SPI Sampler Power Metrics',
    tooltip:
      'Avg/Max power consumption of different components/rails, including max of moving average of window size of 100us/1ms/10ms/1s.',
    description:
      'HBM: HBM SPI Sampler Power Meter; Core: VDD Core SPI Sampler Power Meter',
    getChildValues: (props) =>
      (props as GeneralProps)['spi_power_metrics']?.split('##') || [],
  },
];

/** A performance summary view component. */
@Component({
  changeDetection: ChangeDetectionStrategy.Default,
  standalone: false,
  selector: 'performance-summary',
  templateUrl: './performance_summary.ng.html',
  styleUrls: ['./performance_summary.scss'],
})
export class PerformanceSummary implements OnChanges, OnInit {
  private readonly dataService = inject(DATA_SERVICE_INTERFACE_TOKEN);
  /** Identify if this is an inference or training session */
  @Input() isInference?: boolean;

  /** The general analysis data. */
  @Input() generalAnalysis?: GeneralAnalysis;

  /** The input pipeline analyis data. */
  @Input() inputPipelineAnalysis?: InputPipelineAnalysis;

  /** Inference latency analysis data */
  @Input() inferenceLatencyData?: GeneralAnalysis;

  @Input() sessionId = '';

  /** Disaggregated serving latency data */
  @Input() disaggregatedServingLatencyData?: GeneralAnalysis;
  @Input() baselineGeneralAnalysis?: GeneralAnalysis;
  @Input() baselineInputPipelineAnalysis?: InputPipelineAnalysis;
  @Input() baselineInferenceLatencyData?: SimpleDataTable;
  @Input() baselineDisaggregatedServingLatencyData?: SimpleDataTable;

  title = 'Performance Summary';
  summaryInfoCombined: SummaryInfo[] = [];

  get isTpu() {
    return this.inputPipelineProps['hardware_type'] === 'TPU';
  }

  get hasBaseline(): boolean {
    return !!this.baselineGeneralAnalysis ||
        !!this.baselineInputPipelineAnalysis ||
        !!this.baselineInferenceLatencyData ||
        !!this.baselineDisaggregatedServingLatencyData;
  }

  get generalProps() {
    return ((this.generalAnalysis || {}).p as GeneralProps) || {};
  }

  get inputPipelineProps() {
    return ((this.inputPipelineAnalysis || {}).p as GeneralProps) || {};
  }

  get inferenceLatencyProps() {
    return ((this.inferenceLatencyData || {}).p as GeneralProps) || {};
  }

  get disaggregatedServingLatencyProps() {
    return (
      ((this.disaggregatedServingLatencyData || {}).p as GeneralProps) || {}
    );
  }

  get baselineGeneralProps() {
    return this.baselineGeneralAnalysis
      ? (this.baselineGeneralAnalysis.p as GeneralProps)
      : undefined;
  }

  get baselineInputPipelineProps() {
    return this.baselineInputPipelineAnalysis
      ? (this.baselineInputPipelineAnalysis.p as GeneralProps)
      : undefined;
  }

  get baselineInferenceLatencyProps() {
    return this.baselineInferenceLatencyData
      ? (this.baselineInferenceLatencyData.p as GeneralProps)
      : undefined;
  }

  get baselineDisaggregatedServingLatencyProps() {
    return this.baselineDisaggregatedServingLatencyData
      ? (this.baselineDisaggregatedServingLatencyData.p as GeneralProps)
      : undefined;
  }

  get remarkText() {
    return this.generalProps['remark_text'] || '';
  }

  get remarkColor() {
    return this.generalProps['remark_color'] || '';
  }

  ngOnInit() {
    this.parseSummaryData();
  }

  ngOnChanges(changes: SimpleChanges) {
    this.parseSummaryData();
  }

  parseSummaryData() {
    this.summaryInfoCombined = [];
    this.parseDataFromConfig(
      GENERIC_SUMMARY_INFO_BEFORE,
      this.inputPipelineProps,
      this.summaryInfoCombined,
      {
        level: 1,
        baselineProps: this.baselineInputPipelineProps,
      },
    );
    if (this.isTpu) {
      this.parseDataFromConfig(
        TPU_SUMMARY_INFO,
        this.generalProps,
        this.summaryInfoCombined,
        {
          level: 1,
          baselineProps: this.baselineGeneralProps,
        },
      );
    } else {
      this.parseDataFromConfig(
        NON_TPU_SUMMARY_INFO,
        this.inputPipelineProps,
        this.summaryInfoCombined,
        {
          level: 1,
          baselineProps: this.baselineInputPipelineProps,
        },
      );
    }
    this.parseDataFromConfig(
      GENERIC_SUMMARY_INFO_AFTER,
      this.generalProps,
      this.summaryInfoCombined,
      {
        level: 1,
        baselineProps: this.baselineGeneralProps,
      },
    );
  }

  readSummaryInfoFromConfig(
    config: SummaryInfoConfig,
    props: GeneralProps,
    options: ParseConfigOptions = {},
  ): SummaryInfo | null {
    if (config.trainingOnly && this.isInference) return null;
    if (config.inferenceOnly && !this.isInference) return null;

    const {customInput, baselineProps, baselineCustomInput, parentGoodMetric} =
      options;

    const descriptions = [];
    // We've seen 'nan' in the sdv value
    if (
      config.sdvKey &&
      props[config.sdvKey] &&
      props[config.sdvKey] !== 'nan'
    ) {
      descriptions.push(
        `(σ = ${props[config.sdvKey] || ''} ${config.unit || ''})`,
      );
    }
    if (config.goodMetric !== undefined) {
      descriptions.push(`${config.goodMetric ? 'higher' : 'lower'} is better.`);
    }
    if (config.description) {
      descriptions.push(config.description);
    }
    let value = '';
    let valueStr = '';
    // valueKey and getValue are mutual exclusive
    if (config.valueKey) {
      value = props[config.valueKey];
      valueStr = config.unit ? `${value} ${config.unit}` : value;
    } else if (config.getValue && customInput) {
      value = config.getValue(customInput);
      valueStr = config.unit ? `${value} ${config.unit}` : value;
    }

    let baselineValue = '';
    let baselineValueStr = '';
    if (
      config.valueKey &&
      baselineProps &&
      baselineProps[config.valueKey] !== undefined
    ) {
      baselineValue = baselineProps[config.valueKey];
      baselineValueStr = config.unit
        ? `${baselineValue} ${config.unit}`
        : baselineValue;
    } else if (
      config.getValue &&
      baselineCustomInput &&
      baselineCustomInput.length > 0
    ) {
      baselineValue = config.getValue(baselineCustomInput);
      baselineValueStr = config.unit
        ? `${baselineValue} ${config.unit}`
        : baselineValue;
    }
    // Add dynamic Roofline link with preserved parameters
    if (config.valueKey === 'flop_rate_utilization_relative_to_roofline') {
      const url = this.dataService.getRooflineModelLink(this.sessionId);
      descriptions.push(`see <a href="${url}">roofline_model</a>`);
    }

    const childrenInfoCombined: SummaryInfo[] = [];
    if (config.childrenInfoConfig) {
      this.parseDataFromConfig(
        config.childrenInfoConfig,
        props,
        childrenInfoCombined,
        {
          level: 0,
          baselineProps,
          baselineCustomInput,
          parentGoodMetric:
            config.goodMetric !== undefined
              ? config.goodMetric
              : parentGoodMetric,
        },
      );
    }

    let diffValueStr = '';
    let diffColor = '';

    if (value && baselineValue) {
      const activeNum = parseToNumber(value);
      const baselineNum = parseToNumber(baselineValue);
      if (activeNum !== null && baselineNum !== null && baselineNum !== 0) {
        const pctDiff = ((activeNum - baselineNum) / baselineNum) * 100;
        const sign = pctDiff > 0 ? '+' : '';
        diffValueStr = `${sign}${pctDiff.toFixed(1)}%`;

          const effectiveGoodMetric =
            config.goodMetric !== undefined
              ? config.goodMetric
              : parentGoodMetric;
          if (effectiveGoodMetric !== undefined) {
            const isBetter = effectiveGoodMetric ? pctDiff > 0 : pctDiff < 0;
            const isWorse = effectiveGoodMetric ? pctDiff < 0 : pctDiff > 0;
            if (isBetter) {
              diffColor = 'green';
            } else if (isWorse) {
              diffColor = 'red';
            }
          }
        }
      }

    const propertyValues = config.getChildValues
      ? config.getChildValues(customInput || props)
      : [];

    if (value || childrenInfoCombined.length > 0 || propertyValues.length > 0) {
      return {
        title: config.title,
        tooltip: config.tooltip || '',
        value: valueStr,
        propertyValues,
        descriptions,
        valueColor: config.valueColor || '',
        childrenInfo: childrenInfoCombined,
        baselineValue: baselineValueStr || undefined,
        diffValue: diffValueStr || undefined,
        diffColor: diffColor || undefined,
      };
    }
    return null;
  }

  // TODO: Remove customInput argument, read all metrics from property
  parseDataFromConfig(
    summaryConfigs: SummaryInfoConfig[] | undefined,
    props: GeneralProps,
    summaryInfoCombined: SummaryInfo[],
    options: ParseConfigOptions = {},
  ) {
    if (!summaryConfigs) return;
    const {level = 0} = options;
    summaryConfigs.forEach((config: SummaryInfoConfig) => {
      const summaryInfo = this.readSummaryInfoFromConfig(
        config,
        props,
        options,
      );
      if (summaryInfo) {
        summaryInfoCombined.push({...summaryInfo, level});
      }
    });
  }
}
