#include <inttypes.h>
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <smmintrin.h>
#include <xmmintrin.h>
#include <immintrin.h>


const uint8_t INF_BYTES = 26;

bool read_pack(int n, FILE* f_in, FILE* f_out)
{
    uint8_t* a = (uint8_t*)_mm_malloc(n * sizeof(uint8_t), sizeof(uint8_t));
    int readed = fread(a, sizeof(*a), n, f_in);
    if (f_out != NULL) {
        fwrite(a, sizeof(*a), readed, f_out);
    }
    _mm_free(a);
    return readed == n;
}

uint32_t read_information(FILE* fin,
    FILE* fout,
    int32_t* width,
    int32_t* height)
{
    read_pack(10, fin, fout);
    uint32_t bOffBits;
    fread(&bOffBits, sizeof(uint32_t), 1, fin);

    if (fout != NULL) {
        fwrite(&bOffBits, sizeof(bOffBits), 1, fout);
    }
    read_pack(4, fin, fout);


    fread(width, sizeof(*width), 1, fin);
    fread(height, sizeof(*height), 1, fin);
    if (fout != NULL) {
        fwrite(width, sizeof(*width), 1, fout);
        fwrite(height, sizeof(*height), 1, fout);
    }
    return bOffBits;
}

void read_from_file(
    FILE* fin,
    FILE* fout,
    uint8_t*** color,
    int32_t* width,
    int32_t* height)
{
    uint32_t bOffBits = read_information(fin, fout, width, height);

    read_pack(bOffBits - INF_BYTES, fin, fout);

    (*color)[0] = (uint8_t*)_mm_malloc((*height) * (*width) * sizeof(uint8_t), sizeof(uint8_t));
    (*color)[1] = (uint8_t*)_mm_malloc((*height) * (*width) * sizeof(uint8_t), sizeof(uint8_t));
    (*color)[2] = (uint8_t*)_mm_malloc((*height) * (*width) * sizeof(uint8_t), sizeof(uint8_t));
    (*color)[3] = (uint8_t*)_mm_malloc((*height) * (*width) * sizeof(uint8_t), sizeof(uint8_t));

    for (int32_t i = 0; i < (*height) * (*width); ++i) {
        fread((*color)[0] + i, sizeof(uint8_t), 1, fin);
        fread((*color)[1] + i, sizeof(uint8_t), 1, fin);
        fread((*color)[2] + i, sizeof(uint8_t), 1, fin);
        fread((*color)[3] + i, sizeof(uint8_t), 1, fin);
    }

}

void recalculate(uint8_t** color1, int32_t width, int32_t height, uint8_t** color2)
{
    uint8_t* recalc = _mm_malloc(sizeof(uint8_t) * 4, sizeof(uint8_t));

    for (int32_t i = 0; i < height; ++i) {
        for (int32_t j = 0; j + 4 < width; j += 4) {
            int32_t ind = i * width + j;
               color1[0][ind] = 255;
               color1[0][ind + 1] = 255;
               color1[0][ind + 2] = 255;
               color1[0][ind + 3] = 255;


               __m256 A2 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color2[0][ind + 3], color2[0][ind + 2], color2[0][ind + 1], color2[0][ind + 0]));
               __m256 A2_255 = _mm256_cvtepi64_pd(_mm256_set_epi64x(
                          255 - color2[0][ind + 3], 255 - color2[0][ind + 2], 255 - color2[0][ind + 1], 255 - color2[0][ind]));
               __m256 max = _mm256_cvtepi64_pd(_mm256_set_epi64x(255, 255, 255, 255));

               __m256 col1 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color1[2][ind + 3], color1[2][ind + 2], color1[2][ind + 1], color1[2][ind + 0]));
               __m256 col2 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color2[2][ind + 3], color2[2][ind + 2], color2[2][ind + 1], color2[2][ind + 0]));

               col1 = _mm256_div_pd(_mm256_add_pd(_mm256_mul_pd(col1, A2_255), _mm256_mul_pd(col2, A2)), max);
               _mm256_store_pd(recalc, col1);
               color1[2][ind] = (uint8_t)recalc[0];
               color1[2][ind + 1] = (uint8_t)recalc[1];
               color1[2][ind + 2] = (uint8_t)recalc[2];
               color1[2][ind + 3] = (uint8_t)recalc[3];



               col1 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color1[1][ind + 3], color1[1][ind + 2], color1[1][ind + 1], color1[1][ind + 0]));
               col2 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color2[1][ind + 3], color2[1][ind + 2], color2[1][ind + 1], color2[1][ind + 0]));

               col1 = _mm256_div_pd(_mm256_add_pd(_mm256_mul_pd(col1, A2_255), _mm256_mul_pd(col2, A2)), max);
               _mm256_store_pd(recalc, col1);
               color1[1][ind] = (uint8_t)recalc[0];
               color1[1][ind + 1] = (uint8_t)recalc[1];
               color1[1][ind + 2] = (uint8_t)recalc[2];
               color1[1][ind + 3] = (uint8_t)recalc[3];

               col1 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color1[3][ind + 3], color1[3][ind + 2], color1[3][ind + 1], color1[3][ind + 0]));
               col2 = _mm256_cvtepi64_pd(
                          _mm256_set_epi64x(color2[3][ind + 3], color2[3][ind + 2], color2[3][ind + 1], color2[3][ind + 0]));
               col1 = _mm256_div_pd(_mm256_add_pd(_mm256_mul_pd(col1, A2_255), _mm256_mul_pd(col2, A2)), max);
               _mm256_store_pd(recalc, col1);
               color1[3][ind] = (uint8_t)recalc[0];
               color1[3][ind + 1] = (uint8_t)recalc[1];
               color1[3][ind + 2] = (uint8_t)recalc[2];
               color1[3][ind + 3] = (uint8_t)recalc[3];
        }
    }
    _mm_free(recalc);
}

void write_to_file(FILE* file, uint32_t height, uint32_t width, uint8_t** color1)
{
    for (int32_t i = 0; i < height * width; ++i) {
        fwrite(color1[0] + i, sizeof(uint8_t), 1, file);
        fwrite(color1[1] + i, sizeof(uint8_t), 1, file);
        fwrite(color1[2] + i, sizeof(uint8_t), 1, file);
        fwrite(color1[3] + i, sizeof(uint8_t), 1, file);
    }
}

void free_arr(uint8_t*** col)
{
    _mm_free((*col)[0]);
    _mm_free((*col)[1]);
    _mm_free((*col)[2]);
    _mm_free((*col)[3]);
    _mm_free(*col);
}

int main(int argc, char* argv[])
{
    FILE* f1 = fopen(argv[1], "rb");
    FILE* f2 = fopen(argv[2], "rb");
    FILE* fout = fopen(argv[3], "wb");


    uint8_t** color1 = (uint8_t**)_mm_malloc(4 * sizeof(uint8_t*), sizeof(uint8_t*));
    int32_t width1;
    int32_t height1;

    uint8_t** color2 = (uint8_t**)_mm_malloc(4 * sizeof(uint8_t*), sizeof(uint8_t*));
    int32_t width2;
    int32_t height2;

    read_from_file(f1, fout, &color1, &width1, &height1);
    read_from_file(f2, NULL, &color2, &width2, &height2);


    recalculate(color1, width1, height1, color2);

    write_to_file(fout, height1, width1, color1);

    while (read_pack(1000, f1, fout)) {
        continue;
    }

    free_arr(&color1);
    free_arr(&color2);
    return 0;
}