import sys
import numpy as np
import scipy
import scipy.sparse
import scipy.sparse.linalg
from imageio import imread, imsave
from skimage.transform import resize as imresize
from skimage.color import rgb2gray
from PIL import Image

import math
from collections import defaultdict

from bresenham import *


def image(filename, size):
    """Original image loading function - kept for backward compatibility"""
    img = imresize(rgb2gray(imread(filename)), (size, size))
    return img


def image_from_pil(pil_image, size):
    """New function to handle PIL images for Gradio compatibility"""
    # Convert PIL image to grayscale if it's not already
    if pil_image.mode != 'L':
        img = pil_image.convert('L')
    else:
        img = pil_image
    
    # Resize image
    img = img.resize((size, size), Image.Resampling.LANCZOS)
    
    # Convert to numpy array and normalize to 0-1 range
    img_array = np.array(img, dtype=np.float64) / 255.0
    
    return img_array


def image_from_array(img_array, size):
    """Handle numpy arrays directly"""
    if len(img_array.shape) == 3:
        # Convert RGB to grayscale if needed
        img_array = rgb2gray(img_array)
    
    # Resize if needed
    if img_array.shape != (size, size):
        img_array = imresize(img_array, (size, size))
    
    return img_array


def build_arc_adjecency_matrix(n, radius):
    print("building sparse adjecency matrix")
    hooks = np.array([[math.cos(np.pi*2*i/n), math.sin(np.pi*2*i/n)] for i in range(n)])
    hooks = (radius * hooks).astype(int)
    edge_codes = []
    row_ind = []
    col_ind = []
    for i, ni in enumerate(hooks):
        for j, nj in enumerate(hooks[i+1:], start=i+1):
            edge_codes.append((i, j))
            pixels = bresenham(ni, nj).path
            edge = []
            for pixel in pixels:
                pixel_code = (pixel[1]+radius)*(radius*2+1) + (pixel[0]+radius)
                edge.append(pixel_code)
            row_ind += edge
            col_ind += [len(edge_codes)-1] * len(edge)
    # creating the edge-pixel adjecency matrix:
    # rows are indexed with pixel codes, columns are indexed with edge codes.
    sparse = scipy.sparse.csr_matrix(([1.0]*len(row_ind), (row_ind, col_ind)), shape=((2*radius+1)*(2*radius+1), len(edge_codes)))
    return sparse, hooks, edge_codes


def build_circle_adjecency_matrix(radius, small_radius):
    print("building sparse adjecency matrix")
    edge_codes = []
    row_ind = []
    col_ind = []
    pixels = circle(small_radius)
    for i, cx in enumerate(range(-radius+small_radius+1, radius-small_radius-1, 1)):
        for j, cy in enumerate(range(-radius+small_radius+1, radius-small_radius-1, 1)):
            edge_codes.append((i, j))
            edge = []
            for pixel in pixels:
                px, py = cx+pixel[0], cy+pixel[1]
                pixel_code = (py+radius)*(radius*2+1) + (px+radius)
                edge.append(pixel_code)
            row_ind += edge
            col_ind += [len(edge_codes)-1] * len(edge)
    # creating the edge-pixel adjecency matrix:
    # rows are indexed with pixel codes, columns are indexed with edge codes.
    sparse = scipy.sparse.csr_matrix(([1.0]*len(row_ind), (row_ind, col_ind)), shape=((2*radius+1)*(2*radius+1), len(edge_codes)))
    hooks = []
    return sparse, hooks, edge_codes


def build_image_vector(img, radius):
    # representing the input image as a sparse column vector of pixels:
    assert img.shape[0] == img.shape[1]
    img_size = img.shape[0]
    row_ind = []
    col_ind = []
    data = []
    for y, line in enumerate(img):
        for x, pixel_value in enumerate(line):
            global_x = x - img_size//2
            global_y = y - img_size//2
            pixel_code = (global_y+radius)*(radius*2+1) + (global_x+radius)
            data.append(float(pixel_value))
            row_ind.append(pixel_code)
            col_ind.append(0)
    sparse_b = scipy.sparse.csr_matrix((data, (row_ind, col_ind)), shape=((2*radius+1)*(2*radius+1), 1))
    return sparse_b


def reconstruct(x, sparse, radius):
    b_approx = sparse.dot(x)
    b_approx = b_approx.toarray().flatten() if scipy.sparse.issparse(b_approx) else b_approx.flatten()
    b_image = b_approx.reshape((2*radius+1, 2*radius+1))
    b_image = np.clip(b_image, 0, 1)  # Changed from 0-255 to 0-1 range
    return b_image


def reconstruct_and_save(x, sparse, radius, filename):
    brightness_correction = 1.2
    b_image = reconstruct(x * brightness_correction, sparse, radius)
    # Convert back to 0-255 range for saving
    b_image_255 = (b_image * 255).astype(np.uint8)
    imsave(filename, b_image_255)


def reconstruct_as_pil(x, sparse, radius, brightness_correction=1.2):
    """New function to return PIL Image directly for Gradio"""
    b_image = reconstruct(x * brightness_correction, sparse, radius)
    # Convert to 0-255 range and create PIL Image
    b_image_255 = (b_image * 255).astype(np.uint8)
    return Image.fromarray(b_image_255, mode='L')


def dump_arcs(solution, hooks, edge_codes, filename):
    f = open(filename, "w")
    n = len(hooks)
    print(n, file=f)
    for i, (x, y) in enumerate(hooks):
        print("%d\t%f\t%f" % (i, x, y), file=f)
    print(file=f)
    assert len(edge_codes) == len(solution)
    for (i, j), value in zip(edge_codes, solution):
        if value==0:
            continue
        # int values are shown as ints.
        if value==int(value):
            value = int(value)
        print("%d\t%d\t%s" % (i, j, str(value)), file=f)
    f.close()


def process_string_art_from_pil(pil_image, n=180, radius=250, quantization_level=30, clip_factor=0.3):
    """
    Main processing function for Gradio - takes PIL image and returns results
    
    Args:
        pil_image: PIL Image object
        n: number of hooks around the circle
        radius: radius of the circle in pixels
        quantization_level: discretization level for string weights
        clip_factor: factor to clip maximum values
    
    Returns:
        tuple: (final_pil_image, stats_dict)
    """
    try:
        # Build adjacency matrix
        sparse, hooks, edge_codes = build_arc_adjecency_matrix(n, radius)
        
        # Process image
        shrinkage = 0.75
        img = image_from_pil(pil_image, int(radius * 2 * shrinkage))
        sparse_b = build_image_vector(img, radius)
        
        # Solve linear system
        print("solving linear system")
        result = scipy.sparse.linalg.lsqr(sparse, np.array(sparse_b.todense()).flatten())
        print("done")
        x = result[0]
        
        # Clip negative values (physically unrealistic)
        x = np.clip(x, 0, 1e6)
        
        # Apply quantization
        max_edge_weight_orig = np.max(x)
        if max_edge_weight_orig > 0:  # Avoid division by zero
            x_quantized = (x / max_edge_weight_orig * quantization_level).round()
            x_quantized = np.clip(x_quantized, 0, int(np.max(x_quantized) * clip_factor))
            # Scale back
            x = x_quantized / quantization_level * max_edge_weight_orig
        else:
            x_quantized = np.zeros_like(x)
        
        # Create final image
        final_pil = reconstruct_as_pil(x, sparse, radius)
        
        # Calculate statistics
        arc_count = int(np.sum(x_quantized)) if len(x_quantized) > 0 else 0
        unique_arcs = len(x_quantized[x_quantized > 0]) if len(x_quantized) > 0 else 0
        
        # Calculate total distance
        total_distance = 0.0
        for edge_code, multiplicity in enumerate(x_quantized):
            if multiplicity > 0:
                hook_index1, hook_index2 = edge_codes[edge_code]
                hook1, hook2 = hooks[hook_index1], hooks[hook_index2]
                distance = np.linalg.norm(hook1.astype(float) - hook2.astype(float)) / radius
                total_distance += distance * multiplicity
        
        stats = {
            'arc_count': arc_count,
            'unique_arcs': unique_arcs,
            'total_distance': total_distance / 2,  # unit diameter, not radius
            'hooks': n,
            'radius': radius,
            'quantization': quantization_level
        }
        
        return final_pil, stats
        
    except Exception as e:
        print(f"Error in processing: {str(e)}")
        # Return a blank image and error stats
        blank_image = Image.new('L', (2*radius+1, 2*radius+1), 128)
        error_stats = {
            'arc_count': 0,
            'unique_arcs': 0,
            'total_distance': 0,
            'hooks': n,
            'radius': radius,
            'quantization': quantization_level,
            'error': str(e)
        }
        return blank_image, error_stats


def main():
    """Original main function - kept for backward compatibility"""
    filename, output_prefix = sys.argv[1:]

    n = 180
    radius = 250

    sparse, hooks, edge_codes = build_arc_adjecency_matrix(n, radius)
    # sparse, hooks, edge_codes = build_circle_adjecency_matrix(radius, 10)

    # square image with same center as the circle, sides are 75% of circle diameter.
    shrinkage = 0.75
    img = image(filename, int(radius * 2 * shrinkage))
    sparse_b = build_image_vector(img, radius)
    # imsave(output_prefix+"-original.png", sparse_b.todense().reshape((2*radius+1, 2*radius+1)))

    # finding the solution, a weighting of edges:
    print("solving linear system")
    # note the .todense(). for some reason the sparse version did not work.
    result = scipy.sparse.linalg.lsqr(sparse, np.array(sparse_b.todense()).flatten())
    print("done")
    # x, istop, itn, r1norm, r2norm, anorm, acond, arnorm = result
    x = result[0]

    reconstruct_and_save(x, sparse, radius, output_prefix+"-allow-negative.png")

    # negative values are clipped, they are physically unrealistic.
    x = np.clip(x, 0, 1e6)

    reconstruct_and_save(x, sparse, radius, output_prefix+"-unquantized.png")
    dump_arcs(x, hooks, edge_codes, output_prefix+"-unquantized.txt")

    # quantizing:
    quantization_level = 30 # 50 is already quite good. None means no quantization.
    # clip values larger than clip_factor times maximum.
    # (The long tail does not add too much to percieved quality.)
    clip_factor = 0.3
    if quantization_level is not None:
        max_edge_weight_orig = np.max(x)
        x_quantized = (x / np.max(x) * quantization_level).round()
        x_quantized = np.clip(x_quantized, 0, int(np.max(x_quantized) * clip_factor))
        # scale it back:
        x = x_quantized / quantization_level * max_edge_weight_orig
        dump_arcs(x_quantized, hooks, edge_codes, output_prefix+".txt")

    reconstruct_and_save(x, sparse, radius, output_prefix+".png")

    if quantization_level is not None:
        arc_count = 0
        total_distance = 0.0
        hist = defaultdict(int)
        for edge_code, multiplicity in enumerate(x_quantized):
            multiplicity = int(multiplicity)
            hist[multiplicity] += 1
            arc_count += multiplicity
            hook_index1, hook_index2 = edge_codes[edge_code]
            hook1, hook2 = hooks[hook_index1], hooks[hook_index2]
            distance = np.linalg.norm(hook1.astype(float) - hook2.astype(float)) / radius
            total_distance += distance * multiplicity
        for multiplicity in range(max(hist.keys())+1):
            print(multiplicity, hist[multiplicity])
        print("total arc count", arc_count)
        print("number of different arcs used", len(x_quantized[x_quantized>0]))
        print("total distance (assuming a unit diameter circle)", total_distance / 2) # unit diameter, not unit radius.


if __name__ == "__main__":
    main()