import { assert, unreachable } from '../../../../../../common/util/util.js';
import { kValue } from '../../../../../util/constants.js';
import {
  Type,
  TypeF16,
  Value,
  elementType,
  elementsOf,
  isAbstractType,
} from '../../../../../util/conversion.js';
import {
  scalarF16Range,
  scalarF32Range,
  scalarF64Range,
  linearRange,
} from '../../../../../util/math.js';
import { ShaderValidationTest } from '../../../shader_validation_test.js';

/// A linear sweep between -2 to 2
export const kMinusTwoToTwo = linearRange(-2, 2, 10);

/// An array of values ranging from -3π to 3π, with a focus on multiples of π
export const kMinus3PiTo3Pi = [
  -3 * Math.PI,
  -2.999 * Math.PI,

  -2.501 * Math.PI,
  -2.5 * Math.PI,
  -2.499 * Math.PI,

  -2.001 * Math.PI,
  -2.0 * Math.PI,
  -1.999 * Math.PI,

  -1.501 * Math.PI,
  -1.5 * Math.PI,
  -1.499 * Math.PI,

  -1.001 * Math.PI,
  -1.0 * Math.PI,
  -0.999 * Math.PI,

  -0.501 * Math.PI,
  -0.5 * Math.PI,
  -0.499 * Math.PI,

  -0.001,
  0,
  0.001,

  0.499 * Math.PI,
  0.5 * Math.PI,
  0.501 * Math.PI,

  0.999 * Math.PI,
  1.0 * Math.PI,
  1.001 * Math.PI,

  1.499 * Math.PI,
  1.5 * Math.PI,
  1.501 * Math.PI,

  1.999 * Math.PI,
  2.0 * Math.PI,
  2.001 * Math.PI,

  2.499 * Math.PI,
  2.5 * Math.PI,
  2.501 * Math.PI,

  2.999 * Math.PI,
  3 * Math.PI,
] as const;

/// A minimal array of values ranging from -3π to 3π, with a focus on multiples
/// of π. Used when multiple parameters are being passed in, so the number of
/// cases becomes the square or more of this list.
export const kSparseMinus3PiTo3Pi = [
  -3 * Math.PI,
  -2.5 * Math.PI,
  -2.0 * Math.PI,
  -1.5 * Math.PI,
  -1.0 * Math.PI,
  -0.5 * Math.PI,
  0,
  0.5 * Math.PI,
  Math.PI,
  1.5 * Math.PI,
  2.0 * Math.PI,
  2.5 * Math.PI,
  3 * Math.PI,
] as const;

/// The evaluation stages to test
export const kConstantAndOverrideStages = ['constant', 'override'] as const;

export type ConstantOrOverrideStage = 'constant' | 'override';

/**
 * @returns true if evaluation stage `stage` supports expressions of type @p.
 */
export function stageSupportsType(stage: ConstantOrOverrideStage, type: Type) {
  if (stage === 'override' && isAbstractType(elementType(type)!)) {
    // Abstract numerics are concretized before being used in an override expression.
    return false;
  }
  return true;
}

/**
 * Runs a validation test to check that evaluation of `builtin` either evaluates with or without
 * error at shader creation time or pipeline creation time.
 * @param t the ShaderValidationTest
 * @param builtin the name of the builtin
 * @param expectedResult false if an error is expected, true if no error is expected
 * @param args the arguments to pass to the builtin
 * @param stage the evaluation stage
 */
export function validateConstOrOverrideBuiltinEval(
  t: ShaderValidationTest,
  builtin: string,
  expectedResult: boolean,
  args: Value[],
  stage: ConstantOrOverrideStage
) {
  const elTys = args.map(arg => elementType(arg.type)!);
  const enables = elTys.some(ty => ty === TypeF16) ? 'enable f16;' : '';

  switch (stage) {
    case 'constant': {
      t.expectCompileResult(
        expectedResult,
        `${enables}
const v = ${builtin}(${args.map(arg => arg.wgsl()).join(', ')});`
      );
      break;
    }
    case 'override': {
      assert(!elTys.some(ty => isAbstractType(ty)));
      const constants: Record<string, number> = {};
      const overrideDecls: string[] = [];
      const callArgs: string[] = [];
      let numOverrides = 0;
      for (const arg of args) {
        const argOverrides: string[] = [];
        for (const el of elementsOf(arg)) {
          const name = `o${numOverrides++}`;
          overrideDecls.push(`override ${name} : ${el.type};`);
          argOverrides.push(name);
          constants[name] = Number(el.value);
        }
        callArgs.push(`${arg.type}(${argOverrides.join(', ')})`);
      }
      t.expectPipelineResult({
        expectedResult,
        code: `${enables}
${overrideDecls.join('\n')}
var<private> v = ${builtin}(${callArgs.join(', ')});`,
        constants,
        reference: ['v'],
      });
      break;
    }
  }
}

/** @returns a sweep of the representable values for element type of `type` */
export function fullRangeForType(type: Type, count?: number) {
  if (count === undefined) {
    count = 25;
  }
  switch (elementType(type)?.kind) {
    case 'abstract-float':
      return scalarF64Range({
        pos_sub: Math.ceil((count * 1) / 5),
        pos_norm: Math.ceil((count * 4) / 5),
      });
    case 'f32':
      return scalarF32Range({
        pos_sub: Math.ceil((count * 1) / 5),
        pos_norm: Math.ceil((count * 4) / 5),
      });
    case 'f16':
      return scalarF16Range({
        pos_sub: Math.ceil((count * 1) / 5),
        pos_norm: Math.ceil((count * 4) / 5),
      });
    case 'i32':
      return linearRange(kValue.i32.negative.min, kValue.i32.positive.max, count).map(f =>
        Math.floor(f)
      );
    case 'u32':
      return linearRange(0, kValue.u32.max, count).map(f => Math.floor(f));
  }
  unreachable();
}

/** @returns all the values in the provided arrays with duplicates removed */
export function unique<T>(...arrays: Array<readonly T[]>): T[] {
  const set = new Set<T>();
  for (const arr of arrays) {
    for (const item of arr) {
      set.add(item);
    }
  }
  return [...set];
}
