import os
import sys
import numpy as np
from astropy.io import fits

PRESCAN_X, IMAGE_X, OVERSCAN_X = 27, 1152, 71
PRESCAN_Y, IMAGE_Y, OVERSCAN_Y = 0, 4616, 84
BLOCK_WIDTH = PRESCAN_X + IMAGE_X + OVERSCAN_X


def process_single_file(filename):
    with fits.open(filename) as hdul:
        data = hdul[1].data

    data = data[PRESCAN_Y:PRESCAN_Y + IMAGE_Y, :]
    blocks = [data[:, i * BLOCK_WIDTH + PRESCAN_X:(i + 1) * BLOCK_WIDTH - OVERSCAN_X] for i in range(16)]

    blocks_a = np.concatenate(blocks[:4], axis=1)
    blocks_b = np.concatenate([np.fliplr(b) for b in blocks[4:8]], axis=1)
    blocks_c = np.concatenate([np.flipud(np.fliplr(b)) for b in blocks[11:7:-1]], axis=1)
    blocks_d = np.concatenate([np.flipud(b) for b in blocks[15:11:-1]], axis=1)

    blocks_dc = np.concatenate([blocks_d, blocks_c], axis=1)
    blocks_ab = np.concatenate([blocks_a, blocks_b], axis=1)
    blocks_final = np.concatenate([blocks_ab, blocks_dc], axis=0)

    output_path = os.path.splitext(fn)[0]+'_cropping.fits'
    fits.writeto(output_path, blocks_final, overwrite=True)
    print(f"OK:{output_path}")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        fn = sys.argv[1]
    else:
        fn = 'CSST_MSC_MS_WIDE_20281024132057_20281024132327_10100484833_08_L0_V01.fits'

    process_single_file(fn)