import sys

import numpy
from pywsclean import *

if len(sys.argv) < 2:
    print("Syntax: example.py <ms>\n")
else:
    parameters = ImagingParameters()
    parameters.msPath = sys.argv[1]
    parameters.imageWidth = 1024
    parameters.imageHeight = 1024
    parameters.pixelScaleX = "1amin"
    parameters.pixelScaleY = "1amin"
    parameters.extraParameters = "-weight natural"

    # Test the operator
    with Operator(parameters) as o:
        data, weights = o.read()

        image = numpy.zeros(parameters.imageWidth * parameters.imageHeight)

        o.backward(image, data)

        o.forward(data, image)

        data = numpy.ones(o.data_size(), dtype=numpy.complex128)

        o.backward(image, data)

        o.write("name.fits", image)

        try:
            o.backward(data, image)  # data and image are swapped
        except Exception as e:
            print("- Ok, specifying wrong input raised:", e)

    try:
        o.backward(image, data)  # Outside with block, should raise
    except Exception as e:
        print("- Ok, using operator outside with block raised:", e)

    # Test if a second use of operator works fine
    with Operator(parameters) as o:
        data, weights = o.read()  # Reread data
        try:
            o.backward(
                image, numpy.ones(o.data_size(), dtype=numpy.complex64)
            )  # Use wrong type
        except Exception as e:
            print("- Ok, wrong type in backward raised:", e)
        try:
            o.forward(
                numpy.ones(o.data_size(), dtype=numpy.complex64), image
            )  # Use wrong type
        except Exception as e:
            print("- Ok, wrong type in forward raised:", e)

    # Test the full cleaning command
    wsc = WSClean()
    wsc.width = 1536
    wsc.height = 1536
    wsc.scale = "5asec"
    wsc.datacolumn = "DATA"
    wsc.niter = 250

    wsc.set_uniform_weighting()
    # Alternatively: wsc.set_natural_weighting()

    # This makes wsclean-dirty.fits (et al) from the DATA column
    wsc.image([sys.argv[1]], "wsclean")

    # This predicts from wsclean-model.fits into the MODEL_DATA column
    # (note that the column is always called MODEL_DATA, the datacolumn parameter
    #  only sets the column used in the 'image' command)
    wsc.predict([sys.argv[1]], "wsclean")