js_npz/npz-reader.js

/**
 * Чтение NPZ файлов в браузере без внешних библиотек
 * NPZ файл - это ZIP архив, содержащий .npy файлы
 */

// Проверка поддержки форматов сжатия
const SUPPORTED_DEFLATE_FORMATS = [];
if (typeof DecompressionStream !== 'undefined') {
  // Проверяем поддержку форматов
  // Примечание: создание DecompressionStream может не выбрасывать ошибку сразу,
  // поэтому проверяем при использовании
  // Приоритет: 'deflate-raw' для ZIP (raw deflate), затем 'deflate' (zlib формат)
  SUPPORTED_DEFLATE_FORMATS.push('deflate-raw', 'deflate');
}

class NPZReader {
  /**
   * Читает NPZ файл из ArrayBuffer
   * @param {ArrayBuffer} buffer - ArrayBuffer с содержимым NPZ файла
   * @returns {Promise<Object>} Объект с ключами - именами массивов, значениями - массивами данных
   */
  static async readNPZ(buffer) {
    const zipFiles = await this.parseZIP(buffer);
    console.log(`Найдено файлов в ZIP: ${Object.keys(zipFiles).length}`);
    console.log('Файлы:', Object.keys(zipFiles));
    
    const result = {};
    
    for (const [filename, fileData] of Object.entries(zipFiles)) {
      if (filename.endsWith('.npy')) {
        const arrayName = filename.replace('.npy', '');
        console.log(`Парсинг NPY файла: ${filename} -> ${arrayName}`);
        try {
          result[arrayName] = this.parseNPY(fileData);
          console.log(`NPY файл ${filename} успешно распарсен`);
        } catch (err) {
          console.error(`Ошибка парсинга NPY файла ${filename}:`, err);
          throw err;
        }
      } else {
        console.log(`Пропущен файл (не .npy): ${filename}`);
      }
    }
    
    console.log(`Итого распарсено массивов: ${Object.keys(result).length}`);
    return result;
  }

  /**
   * Парсит ZIP архив (упрощённая реализация для NPZ файлов)
   * @param {ArrayBuffer} buffer - ArrayBuffer с ZIP архивом
   * @returns {Promise<Object>} Объект с именами файлов и их содержимым
   */
  static async parseZIP(buffer) {
    const view = new DataView(buffer);
    const files = {};
    let offset = 0;

    // Ищем центральный каталог в конце файла
    // ZIP файл имеет структуру: [локальные файлы] [центральный каталог] [конец центрального каталога]
    
    // Ищем сигнатуру конца центрального каталога (0x06054b50)
    let eocdOffset = -1;
    for (let i = buffer.byteLength - 22; i >= 0; i--) {
      if (view.getUint32(i, true) === 0x06054b50) {
        eocdOffset = i;
        break;
      }
    }

    if (eocdOffset === -1) {
      throw new Error('Неверный формат ZIP файла');
    }

    // Читаем информацию о центральном каталоге
    const cdOffset = view.getUint32(eocdOffset + 16, true);
    const cdSize = view.getUint32(eocdOffset + 12, true);
    const totalEntries = view.getUint16(eocdOffset + 10, true);

    // Читаем записи центрального каталога
    let currentOffset = cdOffset;
    for (let i = 0; i < totalEntries; i++) {
      // Сигнатура центрального файла (0x02014b50)
      if (view.getUint32(currentOffset, true) !== 0x02014b50) {
        throw new Error('Ошибка чтения центрального каталога');
      }

      // Флаги сжатия
      const flags = view.getUint16(currentOffset + 8, true);
      const method = view.getUint16(currentOffset + 10, true);
      
      // Размеры
      const compressedSize = view.getUint32(currentOffset + 20, true);
      const uncompressedSize = view.getUint32(currentOffset + 24, true);
      
      // Длина имени файла
      const filenameLength = view.getUint16(currentOffset + 28, true);
      const extraLength = view.getUint16(currentOffset + 30, true);
      const commentLength = view.getUint16(currentOffset + 32, true);
      
      // Смещение локального заголовка
      const localHeaderOffset = view.getUint32(currentOffset + 42, true);
      
      // Читаем имя файла
      const filenameBytes = new Uint8Array(buffer, currentOffset + 46, filenameLength);
      const filename = new TextDecoder('utf-8').decode(filenameBytes);
      
      // Отладочная информация
      console.log(`Файл: ${filename}, метод сжатия: ${method}, размер: ${compressedSize}/${uncompressedSize}`);
      
      // Переходим к следующей записи
      currentOffset += 46 + filenameLength + extraLength + commentLength;

      // Читаем локальный заголовок файла
      if (view.getUint32(localHeaderOffset, true) !== 0x04034b50) {
        throw new Error('Ошибка чтения локального заголовка');
      }

      const localFilenameLength = view.getUint16(localHeaderOffset + 26, true);
      const localExtraLength = view.getUint16(localHeaderOffset + 28, true);
      const dataOffset = localHeaderOffset + 30 + localFilenameLength + localExtraLength;

      // Читаем данные файла
      if (method === 0) {
        // Без сжатия (store)
        console.log(`Чтение несжатого файла ${filename} (размер: ${(uncompressedSize / 1024).toFixed(2)} КБ)...`);
        const fileData = new Uint8Array(buffer, dataOffset, uncompressedSize);
        // Создаём копию данных, чтобы иметь независимый ArrayBuffer
        const fileDataCopy = new Uint8Array(fileData);
        files[filename] = fileDataCopy.buffer;
        console.log(`Файл ${filename} прочитан (${fileDataCopy.length} байт)`);
      } else if (method === 8) {
        // Deflate сжатие - используем встроенный API браузера
        // ZIP использует "raw deflate" (без zlib заголовка)
        console.log(`Начало распаковки файла ${filename} (${(compressedSize / 1024 / 1024).toFixed(2)} МБ -> ${(uncompressedSize / 1024 / 1024).toFixed(2)} МБ)...`);
        const compressedData = new Uint8Array(buffer, dataOffset, compressedSize);
        try {
          const startTime = Date.now();
          const decompressedData = await this.inflate(compressedData);
          const endTime = Date.now();
          console.log(`Файл ${filename} распакован за ${((endTime - startTime) / 1000).toFixed(2)} секунд`);
          files[filename] = decompressedData.buffer;
        } catch (err) {
          throw new Error(`Ошибка распаковки файла ${filename}: ${err.message}`);
        }
      } else {
        throw new Error(`Неподдерживаемый метод сжатия: ${method} для файла ${filename}`);
      }
    }

    return files;
  }

  /**
   * Распаковывает данные, сжатые методом Deflate
   * @param {Uint8Array} data - Сжатые данные
   * @returns {Promise<Uint8Array>} Распакованные данные
   */
  static async inflate(data) {
    // Используем встроенный API DecompressionStream (доступен в современных браузерах)
    if (typeof DecompressionStream === 'undefined') {
      throw new Error('DecompressionStream не поддерживается. Используйте современный браузер (Chrome 80+, Firefox 113+, Safari 16.4+).');
    }

    // ZIP использует "raw deflate" (без zlib заголовка и контрольной суммы)
    // Используем поддерживаемые форматы
    const formats = SUPPORTED_DEFLATE_FORMATS.length > 0 
      ? SUPPORTED_DEFLATE_FORMATS 
      : ['deflate-raw', 'deflate']; // Fallback на стандартные форматы
    
    let lastError = null;
    
    for (const format of formats) {
      try {
        console.log(`Попытка распаковки с форматом: ${format}`);
        // Создаём новый stream для каждой попытки
        const stream = new DecompressionStream(format);
        const writer = stream.writable.getWriter();
        const reader = stream.readable.getReader();
        
        // Запускаем чтение параллельно с записью
        const readPromise = (async () => {
          const chunks = [];
          let done = false;
          let totalRead = 0;
          while (!done) {
            const { done: readerDone, value } = await reader.read();
            done = readerDone;
            if (value) {
              chunks.push(value);
              totalRead += value.length;
              // Логируем прогресс каждые 10 МБ
              if (totalRead % (10 * 1024 * 1024) < value.length) {
                console.log(`Распаковано: ${(totalRead / 1024 / 1024).toFixed(2)} МБ`);
              }
            }
          }
          console.log(`Распаковка завершена, всего: ${(totalRead / 1024 / 1024).toFixed(2)} МБ`);
          return { chunks, totalRead };
        })();
        
        // Записываем данные частями, чтобы не блокировать поток
        console.log('Запись данных в stream...');
        const chunkSize = 1024 * 1024; // 1 МБ за раз
        let written = 0;
        for (let i = 0; i < data.length; i += chunkSize) {
          const chunk = data.slice(i, Math.min(i + chunkSize, data.length));
          await writer.write(chunk);
          written += chunk.length;
          if (written % (1024 * 1024 * 5) < chunkSize) { // Логируем каждые 5 МБ
            console.log(`Записано: ${(written / 1024 / 1024).toFixed(2)} МБ из ${(data.length / 1024 / 1024).toFixed(2)} МБ`);
          }
        }
        await writer.close();
        console.log('Данные записаны, ожидание завершения распаковки...');
        
        // Ждём завершения чтения
        const { chunks, totalRead } = await readPromise;
        
        // Объединяем все чанки в один массив
        if (chunks.length === 0) {
          throw new Error('Нет данных после распаковки');
        }
        
        const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
        const result = new Uint8Array(totalLength);
        let offset = 0;
        for (const chunk of chunks) {
          result.set(chunk, offset);
          offset += chunk.length;
        }
        
        return result;
      } catch (err) {
        lastError = err;
        // Пробуем следующий формат
        console.log(`Формат ${format} не сработал:`, err.message);
        continue;
      }
    }
    
    // Если ни один формат не сработал
    throw new Error(`Ошибка распаковки: ${lastError?.message || 'неизвестная ошибка'}. Попробованы форматы: ${formats.join(', ')}. ZIP использует raw deflate, который может не поддерживаться в вашем браузере.`);
  }

  /**
   * Парсит NPY файл (формат NumPy)
   * @param {ArrayBuffer} buffer - ArrayBuffer с содержимым NPY файла
   * @returns {Object} Объект с данными массива {data, shape, dtype}
   */
  static parseNPY(buffer) {
    const view = new DataView(buffer);
    let offset = 0;

    // Проверяем магический заголовок (0x93 'NUMPY')
    if (view.getUint8(offset) !== 0x93) {
      throw new Error('Неверный формат NPY файла: неверный магический байт');
    }
    offset++;

    const magic = String.fromCharCode(
      view.getUint8(offset++),
      view.getUint8(offset++),
      view.getUint8(offset++),
      view.getUint8(offset++),
      view.getUint8(offset++)
    );

    if (magic !== 'NUMPY') {
      throw new Error('Неверный формат NPY файла: неверная магическая строка');
    }

    // Читаем версию формата
    const majorVersion = view.getUint8(offset++);
    const minorVersion = view.getUint8(offset++);

    // Читаем длину заголовка
    let headerLength;
    if (majorVersion === 1) {
      headerLength = view.getUint16(offset, true);
      offset += 2;
    } else if (majorVersion === 2 || majorVersion === 3) {
      headerLength = view.getUint32(offset, true);
      offset += 4;
    } else {
      throw new Error(`Неподдерживаемая версия NPY формата: ${majorVersion}.${minorVersion}`);
    }

    // Читаем заголовок (Python dict в виде строки)
    const headerBytes = new Uint8Array(buffer, offset, headerLength);
    const headerStr = new TextDecoder('utf-8').decode(headerBytes);
    offset += headerLength;

    // Парсим заголовок (упрощённый парсинг Python dict)
    console.log('Заголовок NPY:', headerStr);
    const header = this.parseHeader(headerStr);
    console.log('Распарсенный заголовок:', header);

    // Читаем данные
    const dataOffset = offset;
    const dataLength = buffer.byteLength - dataOffset;
    const data = new Uint8Array(buffer, dataOffset, dataLength);
    
    console.log(`Данные NPY: offset=${dataOffset}, length=${dataLength}, shape=${JSON.stringify(header.shape)}, dtype=${header.descr}`);

    // Проверяем, что shape и descr найдены
    if (!header.descr) {
      throw new Error('Не найден dtype (descr) в заголовке NPY');
    }
    if (!header.shape || header.shape.length === 0) {
      throw new Error('Не найден shape в заголовке NPY');
    }

    // Конвертируем данные в нужный тип
    const typedArray = this.convertToTypedArray(data, header.descr, header.shape);
    console.log(`Конвертировано: ${typedArray.length} элементов`);

    return {
      data: typedArray,
      shape: header.shape,
      dtype: header.descr,
      fortran_order: header.fortran_order || false
    };
  }

  /**
   * Парсит заголовок NPY файла (Python dict)
   * @param {string} headerStr - Строка заголовка
   * @returns {Object} Распарсенный заголовок
   */
  static parseHeader(headerStr) {
    // Упрощённый парсер Python dict
    // Формат: {'descr': '<f8', 'fortran_order': False, 'shape': (2, 3), }
    
    const result = {};

    // Извлекаем dtype (descr)
    const descrMatch = headerStr.match(/descr['"]\s*:\s*['"]([^'"]+)['"]/);
    if (descrMatch) {
      result.descr = descrMatch[1];
    } else {
      console.warn('Не найден descr в заголовке:', headerStr);
    }

    // Извлекаем fortran_order
    const fortranMatch = headerStr.match(/fortran_order['"]\s*:\s*(True|False)/);
    if (fortranMatch) {
      result.fortran_order = fortranMatch[1] === 'True';
    }

    // Извлекаем shape
    // Формат может быть: 'shape': (3,) или 'shape': (3, 4) или 'shape': (3,)
    const shapeMatch = headerStr.match(/shape['"]\s*:\s*\(([^)]*)\)/);
    if (shapeMatch) {
      const shapeStr = shapeMatch[1].trim();
      if (shapeStr) {
        // Разбиваем по запятой и парсим числа
        result.shape = shapeStr.split(',').map(s => {
          const trimmed = s.trim();
          return trimmed ? parseInt(trimmed, 10) : null;
        }).filter(x => x !== null); // Убираем null значения
      } else {
        result.shape = []; // Скаляр
      }
    } else {
      console.warn('Не найден shape в заголовке:', headerStr);
      result.shape = [];
    }

    return result;
  }

  /**
   * Конвертирует байты в типизированный массив
   * @param {Uint8Array} data - Сырые данные
   * @param {string} dtype - Описание типа данных (например, '<f8', '<i4')
   * @param {Array<number>} shape - Форма массива
   * @returns {TypedArray} Типизированный массив
   */
  static convertToTypedArray(data, dtype, shape) {
    // Парсим dtype: '<f8' -> little-endian, float64
    // Формат: [endian][type][size]
    // endian: '<' (little), '>' (big), '=' (native), '|' (not applicable, single byte)
    // type: 'f' (float), 'i' (int), 'u' (uint), 'c' (complex), 'b' (bool)
    // size: количество байт (1, 2, 4, 8)

    console.log(`convertToTypedArray: dtype=${dtype}, shape=${JSON.stringify(shape)}, data.length=${data.length}`);
    
    const endian = dtype[0];
    const type = dtype[1];
    const size = parseInt(dtype.substring(2), 10);
    
    // Проверяем, что size валидный
    if (isNaN(size) || size <= 0) {
      throw new Error(`Неверный размер dtype: ${dtype}`);
    }

    const isLittleEndian = endian === '<' || (endian === '=' && this.isLittleEndian());

    // Вычисляем общее количество элементов
    const totalElements = shape.length === 0 ? 1 : shape.reduce((a, b) => a * b, 1);

    let typedArray;

    switch (type) {
      case 'f': // float
        if (size === 4) {
          typedArray = new Float32Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 8) {
          typedArray = new Float64Array(data.buffer, data.byteOffset, totalElements);
        } else {
          throw new Error(`Неподдерживаемый размер float: ${size}`);
        }
        break;

      case 'i': // int (signed)
        if (size === 1) {
          typedArray = new Int8Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 2) {
          typedArray = new Int16Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 4) {
          typedArray = new Int32Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 8) {
          // Int64Array не существует, используем BigInt64Array
          typedArray = new BigInt64Array(data.buffer, data.byteOffset, totalElements);
        } else {
          throw new Error(`Неподдерживаемый размер int: ${size}`);
        }
        break;

      case 'u': // uint (unsigned)
        if (size === 1) {
          typedArray = new Uint8Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 2) {
          typedArray = new Uint16Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 4) {
          typedArray = new Uint32Array(data.buffer, data.byteOffset, totalElements);
        } else if (size === 8) {
          // Uint64Array не существует, используем BigUint64Array
          typedArray = new BigUint64Array(data.buffer, data.byteOffset, totalElements);
        } else {
          throw new Error(`Неподдерживаемый размер uint: ${size}`);
        }
        break;

      case 'b': // bool
        typedArray = new Uint8Array(data.buffer, data.byteOffset, totalElements);
        break;

      case 'c': // complex
        // Комплексные числа занимают в 2 раза больше места
        if (size === 8) {
          // complex64 = 2 * float32
          typedArray = new Float32Array(data.buffer, data.byteOffset, totalElements * 2);
        } else if (size === 16) {
          // complex128 = 2 * float64
          typedArray = new Float64Array(data.buffer, data.byteOffset, totalElements * 2);
        } else {
          throw new Error(`Неподдерживаемый размер complex: ${size}`);
        }
        break;

      default:
        throw new Error(`Неподдерживаемый тип данных: ${type}`);
    }

    // Если порядок байт не совпадает, нужно переставить байты
    if (!isLittleEndian && size > 1) {
      const view = new DataView(typedArray.buffer, typedArray.byteOffset, typedArray.byteLength);
      for (let i = 0; i < typedArray.length; i++) {
        const offset = i * size;
        // Переставляем байты для каждого элемента
        const bytes = new Uint8Array(size);
        for (let j = 0; j < size; j++) {
          bytes[j] = view.getUint8(offset + j);
        }
        bytes.reverse();
        for (let j = 0; j < size; j++) {
          view.setUint8(offset + j, bytes[j]);
        }
      }
    }

    return typedArray;
  }

  /**
   * Проверяет, является ли система little-endian
   * @returns {boolean}
   */
  static isLittleEndian() {
    const buffer = new ArrayBuffer(2);
    const view = new DataView(buffer);
    view.setUint16(0, 0x0001, true);
    return view.getUint8(0) === 1;
  }
}