/*
 * Copyright 2021 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "stim/io/measure_record_reader.h"

#include <algorithm>

using namespace stim;

// Returns true if keyword is found at current position. Returns false if EOF is found at current
// position. Throws otherwise. Uses next output variable for the next character or EOF.
bool maybe_consume_keyword(FILE *in, const std::string &keyword, int &next) {
    next = getc(in);
    if (next == EOF) {
        return false;
    }

    for (char c : keyword) {
        if (c != next) {
            throw std::runtime_error("Failed to find expected string \"" + keyword + "\"");
        }
        next = getc(in);
    }

    return true;
}

bool stim::read_uint64(FILE *in, uint64_t &value, int &next, bool include_next) {
    if (!include_next) {
        next = getc(in);
    }
    if (!isdigit(next)) {
        return false;
    }

    value = 0;
    while (isdigit(next)) {
        uint64_t prev_value = value;
        value *= 10;
        value += next - '0';
        if (value < prev_value) {
            throw std::runtime_error("Integer value read from file was too big");
        }
        next = getc(in);
    }
    return true;
}

MeasureRecordReader::MeasureRecordReader(size_t num_measurements, size_t num_detectors, size_t num_observables)
    : num_measurements(num_measurements), num_detectors(num_detectors), num_observables(num_observables) {
}

size_t MeasureRecordReader::read_records_into(simd_bit_table &out, bool major_index_is_shot_index, size_t max_shots) {
    if (!major_index_is_shot_index) {
        simd_bit_table buf(out.num_minor_bits_padded(), out.num_major_bits_padded());
        size_t r = read_records_into(buf, true, max_shots);
        buf.transpose_into(out);
        return r;
    }

    size_t num_read = 0;
    max_shots = std::min(max_shots, out.num_major_bits_padded());
    while (num_read < max_shots && start_and_read_entire_record(out[num_read])) {
        num_read++;
    }
    return num_read;
}

std::unique_ptr<MeasureRecordReader> MeasureRecordReader::make(
    FILE *in, SampleFormat input_format, size_t num_measurements, size_t num_detectors, size_t num_observables) {
    switch (input_format) {
        case SAMPLE_FORMAT_01:
            return std::unique_ptr<MeasureRecordReader>(
                new MeasureRecordReaderFormat01(in, num_measurements, num_detectors, num_observables));
        case SAMPLE_FORMAT_B8:
            return std::unique_ptr<MeasureRecordReader>(
                new MeasureRecordReaderFormatB8(in, num_measurements, num_detectors, num_observables));
        case SAMPLE_FORMAT_DETS:
            return std::unique_ptr<MeasureRecordReader>(
                new MeasureRecordReaderFormatDets(in, num_measurements, num_detectors, num_observables));
        case SAMPLE_FORMAT_HITS:
            return std::unique_ptr<MeasureRecordReader>(
                new MeasureRecordReaderFormatHits(in, num_measurements, num_detectors, num_observables));
        case SAMPLE_FORMAT_PTB64:
            throw std::invalid_argument("SAMPLE_FORMAT_PTB64 incompatible with SingleMeasurementRecord");
        case SAMPLE_FORMAT_R8:
            return std::unique_ptr<MeasureRecordReader>(
                new MeasureRecordReaderFormatR8(in, num_measurements, num_detectors, num_observables));
        default:
            throw std::invalid_argument("Sample format not recognized by SingleMeasurementRecord");
    }
}

size_t MeasureRecordReader::read_bits_into_bytes(PointerRange<uint8_t> out_buffer) {
    size_t n = 0;
    for (uint8_t &b : out_buffer) {
        b = 0;
        for (size_t k = 0; k < 8; k++) {
            if (is_end_of_record()) {
                return n;
            }
            b |= uint8_t(read_bit()) << k;
            ++n;
        }
    }
    return n;
}

size_t MeasureRecordReader::bits_per_record() const {
    return num_measurements + num_detectors + num_observables;
}

void MeasureRecordReader::move_obs_in_shots_to_mask_assuming_sorted(SparseShot &shot) {
    if (num_observables > 32) {
        throw std::invalid_argument("More than 32 observables. Can't read into SparseShot struct.");
    }

    size_t nd = num_measurements + num_detectors;
    size_t n = nd + num_observables;
    shot.obs_mask = 0;
    while (!shot.hits.empty()) {
        auto top = shot.hits.back();
        if (top < nd) {
            break;
        }
        if (top >= n) {
            throw std::invalid_argument("Hit index from data is too large.");
        }
        shot.hits.pop_back();
        shot.obs_mask ^= 1 << (top - nd);
    }
}

/// 01 format

MeasureRecordReaderFormat01::MeasureRecordReaderFormat01(
    FILE *in, size_t num_measurements, size_t num_detectors, size_t num_observables)
    : MeasureRecordReader(num_measurements, num_detectors, num_observables),
      in(in),
      payload('\n'),
      position(bits_per_record()) {
}

bool MeasureRecordReaderFormat01::start_and_read_entire_record(simd_bits_range_ref dirty_out_buffer) {
    return start_and_read_entire_record_helper(
        [&](size_t k) {
            dirty_out_buffer[k] = false;
        },
        [&](size_t k) {
            dirty_out_buffer[k] = true;
        });
}

bool MeasureRecordReaderFormat01::start_and_read_entire_record(SparseShot &cleared_out) {
    bool result = start_and_read_entire_record_helper(
        [&](size_t k) {
        },
        [&](size_t k) {
            cleared_out.hits.push_back((uint64_t)k);
        });
    move_obs_in_shots_to_mask_assuming_sorted(cleared_out);
    return result;
}

bool MeasureRecordReaderFormat01::read_bit() {
    if (payload == EOF) {
        throw std::out_of_range("Attempt to read past end-of-file");
    }
    if (payload == '\n' || position >= bits_per_record()) {
        throw std::out_of_range("Attempt to read past end-of-record");
    }
    if (payload != '0' && payload != '1') {
        throw std::runtime_error("Expected '0' or '1' because input format was specified as '01'");
    }

    bool bit = payload == '1';
    payload = getc(in);
    ++position;
    return bit;
}

bool MeasureRecordReaderFormat01::next_record() {
    while (payload != EOF && payload != '\n') {
        payload = getc(in);
        if (position++ > bits_per_record()) {
            throw std::runtime_error(
                "Line was too long for input file in 01 format. Expected " + std::to_string(bits_per_record()) +
                " characters but got " + std::to_string(position));
        }
    }

    return start_record();
}

bool MeasureRecordReaderFormat01::start_record() {
    payload = getc(in);
    position = 0;
    return payload != EOF;
}

bool MeasureRecordReaderFormat01::is_end_of_record() {
    bool payload_ended = (payload == EOF || payload == '\n');
    bool expected_end = position >= bits_per_record();
    if (payload_ended && !expected_end) {
        throw std::invalid_argument("Record data (in 01 format) ended early, before expected length.");
    }
    if (!payload_ended && expected_end) {
        throw std::invalid_argument("Record data (in 01 format) did not end by the expected length.");
    }
    return payload_ended;
}

/// B8 format

MeasureRecordReaderFormatB8::MeasureRecordReaderFormatB8(
    FILE *in, size_t num_measurements, size_t num_detectors, size_t num_observables)
    : MeasureRecordReader(num_measurements, num_detectors, num_observables),
      in(in),
      payload(0),
      bits_available(0),
      position(bits_per_record()) {
}

bool MeasureRecordReaderFormatB8::start_and_read_entire_record(simd_bits_range_ref dirty_out_buffer) {
    return start_and_read_entire_record_helper([&](size_t byte_index, uint8_t byte) {
        dirty_out_buffer.u8[byte_index] = (uint8_t)byte;
    });
}

bool MeasureRecordReaderFormatB8::start_and_read_entire_record(SparseShot &cleared_out) {
    bool result = start_and_read_entire_record_helper([&](size_t byte_index, uint8_t byte) {
        size_t bit_offset = byte_index << 3;
        for (size_t r = 0; r < 8; r++) {
            if (byte & (1 << r)) {
                cleared_out.hits.push_back(bit_offset + r);
            }
        }
    });
    move_obs_in_shots_to_mask_assuming_sorted(cleared_out);
    return result;
}

size_t MeasureRecordReaderFormatB8::read_bits_into_bytes(PointerRange<uint8_t> out_buffer) {
    if (out_buffer.empty()) {
        return 0;
    }
    if (position >= bits_per_record()) {
        return 0;
    }

    if (bits_available & 7) {
        // Partial read from before still had trailing bits.
        return MeasureRecordReader::read_bits_into_bytes(out_buffer);
    }

    size_t total_read = 0;
    if (bits_available) {
        *out_buffer.ptr_start = payload & 0xFF;
        out_buffer.ptr_start++;
        bits_available = 0;
        position += 8;
        total_read += 8;
    }

    size_t n_bits = std::min<size_t>(8 * out_buffer.size(), bits_per_record() - position);
    size_t n_bytes = (n_bits + 7) / 8;
    n_bytes = fread(out_buffer.ptr_start, sizeof(uint8_t), n_bytes, in);
    n_bits = std::min<size_t>(8 * n_bytes, n_bits);
    position += n_bits;
    total_read += n_bits;

    return total_read;
}

bool MeasureRecordReaderFormatB8::read_bit() {
    if (position >= bits_per_record()) {
        throw std::out_of_range("Attempt to read past end-of-record");
    }

    maybe_update_payload();

    if (payload == EOF) {
        throw std::out_of_range("Attempt to read past end-of-file");
    }

    bool b = payload & 1;
    payload >>= 1;
    --bits_available;
    ++position;
    return b;
}

bool MeasureRecordReaderFormatB8::next_record() {
    while (!is_end_of_record()) {
        read_bit();
    }
    return start_record();
}

bool MeasureRecordReaderFormatB8::start_record() {
    position = 0;
    bits_available = 0;
    payload = 0;
    maybe_update_payload();
    return payload != EOF;
}

bool MeasureRecordReaderFormatB8::is_end_of_record() {
    return position >= bits_per_record();
}

void MeasureRecordReaderFormatB8::maybe_update_payload() {
    if (bits_available > 0) {
        return;
    }
    payload = getc(in);
    if (payload != EOF) {
        bits_available = 8;
    }
}

/// Hits format

MeasureRecordReaderFormatHits::MeasureRecordReaderFormatHits(
    FILE *in, size_t num_measurements, size_t num_detectors, size_t num_observables)
    : MeasureRecordReader(num_measurements, num_detectors, num_observables),
      in(in),
      buffer(bits_per_record()),
      position_in_buffer(bits_per_record()) {
}

bool MeasureRecordReaderFormatHits::start_and_read_entire_record(simd_bits_range_ref dirty_out_buffer) {
    dirty_out_buffer.prefix_ref(bits_per_record()).clear();
    return start_and_read_entire_record_helper([&](size_t bit_index) {
        dirty_out_buffer[bit_index] = true;
    });
}

bool MeasureRecordReaderFormatHits::start_and_read_entire_record(SparseShot &cleared_out) {
    size_t m = bits_per_record();
    size_t nmd = num_measurements + num_detectors;
    return start_and_read_entire_record_helper([&](size_t bit_index) {
        if (bit_index >= m) {
            throw std::invalid_argument("hit index is too large.");
        }
        if (bit_index < nmd) {
            cleared_out.hits.push_back(bit_index);
        } else {
            cleared_out.obs_mask ^= 1 << (bit_index - nmd);
        }
    });
}

bool MeasureRecordReaderFormatHits::read_bit() {
    if (position_in_buffer >= bits_per_record()) {
        throw std::invalid_argument("Read past end of buffer.");
    }
    return buffer[position_in_buffer++];
}

bool MeasureRecordReaderFormatHits::next_record() {
    return start_record();
}

bool MeasureRecordReaderFormatHits::start_record() {
    int c = getc(in);
    if (c == EOF) {
        return false;
    }
    buffer.clear();
    position_in_buffer = 0;
    bool is_first = true;
    while (c != '\n') {
        uint64_t value;
        if (!read_uint64(in, value, c, is_first)) {
            throw std::runtime_error(
                "Integer didn't start immediately at start of line or after comma in 'hits' format.");
        }
        if (c != ',' && c != '\n') {
            throw std::runtime_error(
                "'hits' format requires  integers to be followed by a comma or newline, but got a '" +
                std::to_string(c) + "'.");
        }
        if (value >= bits_per_record()) {
            throw std::runtime_error(
                "Bits per record is " + std::to_string(bits_per_record()) + " but got a hit value " +
                std::to_string(value) + ".");
        }
        buffer[(size_t)value] ^= true;
        is_first = false;
    }
    return true;
}

bool MeasureRecordReaderFormatHits::is_end_of_record() {
    return position_in_buffer >= bits_per_record();
}

/// R8 format

MeasureRecordReaderFormatR8::MeasureRecordReaderFormatR8(
    FILE *in, size_t num_measurements, size_t num_detectors, size_t num_observables)
    : MeasureRecordReader(num_measurements, num_detectors, num_observables), in(in) {
}

size_t MeasureRecordReaderFormatR8::read_bits_into_bytes(PointerRange<uint8_t> out_buffer) {
    size_t n = 0;
    for (uint8_t &b : out_buffer) {
        b = 0;
        if (buffered_0s >= 8) {
            position += 8;
            buffered_0s -= 8;
            n += 8;
            continue;
        }
        for (size_t k = 0; k < 8; k++) {
            if (is_end_of_record()) {
                return n;
            }
            b |= uint8_t(read_bit()) << k;
            ++n;
        }
    }
    return n;
}

bool MeasureRecordReaderFormatR8::read_bit() {
    if (!buffered_0s && !buffered_1s) {
        bool read_any = maybe_buffer_data();
        assert(read_any);
    }
    if (buffered_0s) {
        buffered_0s--;
        position++;
        return false;
    } else if (buffered_1s) {
        buffered_1s--;
        position++;
        return true;
    } else {
        throw std::invalid_argument("Read past end-of-record.");
    }
}

bool MeasureRecordReaderFormatR8::start_and_read_entire_record(simd_bits_range_ref dirty_out_buffer) {
    dirty_out_buffer.prefix_ref(bits_per_record()).clear();
    return start_and_read_entire_record_helper([&](size_t bit_index) {
        dirty_out_buffer[bit_index] = 1;
    });
}

bool MeasureRecordReaderFormatR8::start_and_read_entire_record(SparseShot &cleared_out) {
    bool result = start_and_read_entire_record_helper([&](size_t bit_index) {
        cleared_out.hits.push_back(bit_index);
    });
    move_obs_in_shots_to_mask_assuming_sorted(cleared_out);
    return result;
}
bool MeasureRecordReaderFormatR8::next_record() {
    while (!is_end_of_record()) {
        read_bit();
    }
    return start_record();
}

bool MeasureRecordReaderFormatR8::start_record() {
    position = 0;
    have_seen_terminal_1 = false;
    return maybe_buffer_data();
}

bool MeasureRecordReaderFormatR8::is_end_of_record() {
    return position == bits_per_record() && have_seen_terminal_1;
}

bool MeasureRecordReaderFormatR8::maybe_buffer_data() {
    assert(buffered_0s == 0);
    assert(buffered_1s == 0);
    if (is_end_of_record()) {
        throw std::invalid_argument("Attempted to read past end-of-record.");
    }

    // Count zeroes until a one is found.
    int r;
    do {
        r = getc(in);
        if (r == EOF) {
            if (buffered_0s == 0 && position == 0) {
                return false;  // No next record.
            }
            throw std::invalid_argument("r8 data ended on a continuation (a 0xFF byte) which is not allowed.");
        }
        buffered_0s += r;
    } while (r == 0xFF);
    buffered_1s = 1;

    // Check if the 1 is the one at end of data.
    size_t total_data = position + buffered_0s + buffered_1s;
    if (total_data == bits_per_record()) {
        int t = getc(in);
        if (t == EOF) {
            throw std::invalid_argument(
                "r8 data ended too early. "
                "The extracted data ended in a 1, but there was no corresponding 0x00 terminator byte for the expected "
                "'fake encoded 1 just after the end of the data' before the input ended.");
        } else if (t != 0) {
            throw std::invalid_argument(
                "r8 data ended too early. "
                "The extracted data ended in a 1, but there was no corresponding 0x00 terminator byte for the expected "
                "'fake encoded 1 just after the end of the data' before any additional data.");
        }
        have_seen_terminal_1 = true;
    } else if (total_data == bits_per_record() + 1) {
        have_seen_terminal_1 = true;
        buffered_1s = 0;
    } else if (total_data > bits_per_record() + 1) {
        throw std::invalid_argument("r8 data encoded a jump past the expected end of encoded data.");
    }
    return true;
}

/// DETS format

bool MeasureRecordReaderFormatDets::start_and_read_entire_record(simd_bits_range_ref dirty_out_buffer) {
    dirty_out_buffer.prefix_ref(bits_per_record()).clear();
    return start_and_read_entire_record_helper([&](size_t bit_index) {
        dirty_out_buffer[bit_index] = true;
    });
}

bool MeasureRecordReaderFormatDets::start_and_read_entire_record(SparseShot &cleared_out) {
    size_t obs_start = num_measurements + num_detectors;
    return start_and_read_entire_record_helper([&](size_t bit_index) {
        if (bit_index < obs_start) {
            cleared_out.hits.push_back(bit_index);
        } else {
            cleared_out.obs_mask ^= 1 << (bit_index - obs_start);
        }
    });
}

MeasureRecordReaderFormatDets::MeasureRecordReaderFormatDets(
    FILE *in, size_t num_measurements, size_t num_detectors, size_t num_observables)
    : MeasureRecordReader(num_measurements, num_detectors, num_observables),
      in(in),
      buffer(num_measurements + num_detectors + num_observables),
      position_in_buffer(num_measurements + num_detectors + num_observables) {
}

bool MeasureRecordReaderFormatDets::read_bit() {
    if (position_in_buffer >= bits_per_record()) {
        throw std::invalid_argument("Read past end of buffer.");
    }
    return buffer[position_in_buffer++];
}

bool MeasureRecordReaderFormatDets::next_record() {
    return start_record();
}

bool MeasureRecordReaderFormatDets::start_record() {
    int c;
    if (!maybe_consume_keyword(in, "shot", c)) {
        return false;
    }
    buffer.clear();
    position_in_buffer = 0;
    while (true) {
        bool had_spacing = c == ' ';
        while (c == ' ') {
            c = getc(in);
        }
        if (c == '\n' || c == EOF) {
            break;
        }
        if (!had_spacing) {
            throw std::invalid_argument("DETS values must be separated by spaces.");
        }
        size_t offset;
        size_t size;
        char prefix = c;
        if (prefix == 'M') {
            offset = 0;
            size = num_measurements;
        } else if (prefix == 'D') {
            offset = num_measurements;
            size = num_detectors;
        } else if (prefix == 'L') {
            offset = num_measurements + num_detectors;
            size = num_observables;
        } else {
            throw std::invalid_argument("Unrecognized DETS prefix: '" + std::to_string(c) + "'");
        }
        uint64_t number;
        if (!read_uint64(in, number, c)) {
            throw std::invalid_argument("DETS prefix '" + std::to_string(c) + "' wasn't not followed by an integer.");
        }
        if (number >= size) {
            throw std::invalid_argument(
                "Got '" + std::to_string(c) + std::to_string(number) + "' but expected num values of that type is " +
                std::to_string(size) + ".");
        }
        buffer[(size_t)(offset + number)] ^= true;
    }
    return true;
}

bool MeasureRecordReaderFormatDets::is_end_of_record() {
    return position_in_buffer == bits_per_record();
}
