from urllib.request import urlretrieve
from logging import Logger
import math
import os

from PIL import Image

from rollbot import as_command, RollbotFailure, Attachment
from rollbot.injection import Args


@as_command
def seychelles(argstr: Args, logger: Logger):
    if "i.imgur.com" not in argstr:
        RollbotFailure.INVALID_ARGUMENTS.raise_exc(
            detail="The !seychelles command requires a direct imgur link argument."
        )

    try:
        urlretrieve(argstr, "/tmp/seychelles")
    except:
        logger.exception("Image download failed")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Couldn't download the provided image")

    s = Seychelles(logger, "/tmp/seychelles", name_out="/tmp/seychelles_out", ext_out="png")
    s.seychelles(verbose=True)
    s.save()

    with open("/tmp/seychelles_out.png", "rb") as s_out:
        return Attachment(name="image", body=s_out.read())


# By Akshay Chitale for r/vexillology on Reddit
# With small modifications to replace print statements

# For Python 3


class Seychelles:
    def __init__(self, logger, name_in, size_out=None, name_out=None, ext_out=None):
        self.logger = logger
        # Set up input
        self.name_in, self.ext_in = os.path.splitext(name_in)
        self.img_raw = Image.open(name_in)
        self.img_raw = self.img_raw.convert("RGB")
        self.size_in = self.img_raw.size
        # Flip so that colors are measured from top left
        self.img_in = self.img_raw.transpose(Image.FLIP_TOP_BOTTOM)

        # Set up output
        self.size_out = size_out if size_out else self.size_in
        self.name_out = name_out if name_out else self.name_in + "_out"
        self.ext_out = "." + ext_out if ext_out else self.ext_in
        self.img_out = Image.new("RGB", self.size_out)
        self.pixels_out = self.img_out.load()

        # Set up image to print
        self.img_print = None

    def _angle_transfer(self, diagonal, seychelles):
        # Define transfer curve as a parabola
        x1, y1 = 0, 0
        x2, y2 = (diagonal, math.pi / 4)
        x3, y3 = math.pi / 2, math.pi / 2
        # Below code from http://chris35wills.github.io/parabola_python/
        denom = (x1 - x2) * (x1 - x3) * (x2 - x3)
        A = (x3 * (y2 - y1) + x2 * (y1 - y3) + x1 * (y3 - y2)) / denom
        B = (x3 * x3 * (y1 - y2) + x2 * x2 * (y3 - y1) + x1 * x1 * (y2 - y3)) / denom
        # C = 0 guaranteed because parabola goes through (0,0)
        if seychelles:
            # Going forward, use parabola
            return lambda x: A * x * x + B * x
        else:
            # Inverse, solve parabola y=ax^2+bx
            # Watch out for squares, where diagonal is pi/4 so A=0
            if abs(diagonal - math.pi / 4) < 1e-9:
                return lambda x: x
            else:
                return lambda x: (-1 * B + math.sqrt(B * B - 4 * A * (-x))) / (2 * A)

    def seychelles(self, verbose=False):
        # Diagonal angle of output image
        out_diagonal = math.atan2(self.size_out[1], self.size_out[0])
        angle_transfer = self._angle_transfer(out_diagonal, True)
        # Find output color for each output pixel
        if verbose:
            self.logger.info(" Progress:   0%")
        for x in range(self.size_out[0]):
            if verbose:
                self.logger.info(
                    "\r Progress: " + str(int(100 * x / self.size_out[0])).rjust(3) + "%"
                )
            for y in range(self.size_out[1]):
                # First, get the angle
                out_angle = math.atan2(y, x)

                # Then, follow the vector to the end of the flag
                out_x, out_y = x, y
                if x == 0 and y == 0:
                    out_x, out_y = 1.0, 1.0
                elif out_angle < out_diagonal:
                    # Scale by x
                    out_x *= self.size_out[0] * 1.0 / x
                    out_y *= self.size_out[0] * 1.0 / x
                else:
                    # Scale by y
                    out_x *= self.size_out[1] * 1.0 / y
                    out_y *= self.size_out[1] * 1.0 / y

                # Get ratio of point radius to full radius
                point_rad = math.sqrt(x * x + y * y)
                out_rad = math.sqrt(out_x * out_x + out_y * out_y)
                rad_ratio = point_rad / out_rad

                # Coordinates on the input are:
                # x = radius, scaled by width of input
                in_x = rad_ratio * self.size_in[0]
                # y = angle, scaled by height of input
                in_y = angle_transfer(out_angle) * self.size_in[1] * 2.0 / math.pi

                # Ensure coordinates are within range
                in_x_int = int(round(in_x))
                if in_x_int < 0:
                    in_x_int = 0
                elif in_x_int >= self.size_in[0]:
                    in_x_int = self.size_in[0] - 1
                in_y_int = int(round(in_y))
                if in_y_int < 0:
                    in_y_int = 0
                elif in_y_int >= self.size_in[1]:
                    in_y_int = self.size_in[1] - 1

                # Assign input color to output color
                self.pixels_out[x, y] = self.img_in.getpixel((in_x_int, in_y_int))
        if verbose:
            self.logger.info("\r Progress: 100%")
        # Flip so that seychelles is from bottom left
        self.img_print = self.img_out.transpose(Image.FLIP_TOP_BOTTOM)

    def inverse_seychelles(self, verbose=False):
        # Diagonal angle of input image
        in_diagonal = math.atan2(self.size_in[1], self.size_in[0])
        angle_transfer = self._angle_transfer(in_diagonal, False)
        # Find output color for each output pixel
        if verbose:
            self.logger.info(" Progress:   0%")
        for x in range(self.size_out[0]):
            if verbose:
                self.logger.info(
                    "\r Progress: " + str(int(100 * x / self.size_out[0])).rjust(3) + "%"
                )
            for y in range(self.size_out[1]):
                # First, get the angle
                in_angle = angle_transfer(y * math.pi / 2.0 / self.size_out[1])

                # Get ratio of point to full width
                rad_ratio = x * 1.0 / self.size_out[0]

                # Then, follow the vector to the end of the flag
                if in_angle < in_diagonal:
                    # Find by x
                    in_x = self.size_in[0] * 1.0
                    in_y = in_x * math.tan(in_angle)
                else:
                    # Find by y
                    in_y = self.size_in[1] * 1.0
                    in_x = in_y / math.tan(in_angle)
                # Scale by radius ratio
                in_x, in_y = rad_ratio * in_x, rad_ratio * in_y

                # Ensure coordinates are within range
                in_x_int = int(round(in_x))
                if in_x_int < 0:
                    in_x_int = 0
                elif in_x_int >= self.size_in[0]:
                    in_x_int = self.size_in[0] - 1
                in_y_int = int(round(in_y))
                if in_y_int < 0:
                    in_y_int = 0
                elif in_y_int >= self.size_in[1]:
                    in_y_int = self.size_in[1] - 1

                # Assign input color to output color
                self.pixels_out[x, y] = self.img_in.getpixel((in_x_int, in_y_int))
        if verbose:
            self.logger.info("\r Progress: 100%")
        # Flip so that seychelles is from bottom left
        self.img_print = self.img_out.transpose(Image.FLIP_TOP_BOTTOM)

    def save(self, name_out=None, ext_out=None):
        if self.img_print is None:
            raise Exception("No processing done yet")
        name = name_out if name_out else self.name_out
        ext = "." + ext_out if ext_out else self.ext_out
        self.img_print.save(name + ext)

    def show(self):
        if self.img_print is None:
            raise Exception("No processing done yet")
        self.img_print.show()
