# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: all
#     notebook_metadata_filter: all,-language_info,-toc,-latex_envs
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.6.1-dev
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
import copy
import datetime
import pprint
from pathlib import Path

import cartopy
import cartopy.crs
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pytz
import rasterio
from IPython.display import display
from matplotlib import pyplot as plt
from matplotlib.colors import Normalize
from pyproj import CRS, Transformer

# %%
import a301_lib

pacific = pytz.timezone("US/Pacific")
date = datetime.datetime.today().astimezone(pacific)
print(f"written on {date}")

# %% [markdown]
# (assign6b_sol)=
# # Solution: adding an image with features to a cartopy map
#
# Here's a demo of adding features to an image.  The notebook does some extra work to
# illustrate how to make a map that is bigger than the image, and how to reproject from
# one raster crs to another using rasterio.
#
# 1) read in the `vancouver_345_refl.tiff` band 5 image
# 2) find the image corners in utm10 and geodetic lat/lon (for the features)
# 3) put the image on a map that extends 5 km beyond the image
# 4) add the features to the big map
# 5) create a new crs that is laea centered on lat=50 deg N, lon= -120 deg E
# 6) use rasterio.reproject to reproject the utm10 image onto the laea crs
# 7) draw a 50 row by 100 column box on the image

# %% [markdown]
# ## Read in the geotiff
#
# First read in band 5 from the `vancouver_345_refl.tiff` image that was produced by
# {ref}`rasterio_3bands`:

# %%
notebook_dir = a301_lib.data_share.resolve()
print(notebook_dir)
week10_scene = notebook_dir / "vancouver_345_refl.tiff"

# %%
with rasterio.open(week10_scene) as van_raster:
    b5_refl = van_raster.read(3)
    chan3_tags = van_raster.tags(3)
    crs_10 = van_raster.profile["crs"]
    profile=van_raster.profile
    affine_transform = profile["transform"]
    tags = van_raster.tags()
    print(f"\n\n{tags=}, \n\n{chan3_tags=},\n\n {profile=},\n\n {affine_transform=}\n")

# %% [markdown]
# ## Set up the lat/lon UTM10N transform
#
# I'll want to transform back and forth between UTM and lat/lon (since I think in lat/lon), so create
# a transformer object for this.  Copy code from {ref}`image_zoom`

# %%
from pyproj import CRS, Transformer

p_utm10 = crs_10
p_latlon = CRS.from_epsg(4326)
crs_transform = Transformer.from_crs(p_latlon, p_utm10)
print(p_latlon.to_wkt())

# %% [markdown]
# ## Find the corners of the image
#
# I need the upper left and lower right corners to set the image extent for imshow.  This is
# what the affine transform provides.  I'll use my crs_transform to also get the image_extent
# in lon/lat coords to make sure I'm in the right place

# %%
ul_x_utm10, ul_y_utm10 = affine_transform*(0,0)
lr_x_utm10, lr_y_utm10 = affine_transform*(profile['width'],profile['height'])
print(f"{(ul_x_utm10,ul_y_utm10,lr_x_utm10,lr_y_utm10)=}")
image_extent_utm10 = (ul_x_utm10,lr_x_utm10, lr_y_utm10, ul_y_utm10)

# %%
ul_lon,ul_lat = crs_transform.transform(ul_x_utm10,ul_y_utm10,direction='INVERSE')
lr_lon,lr_lat = crs_transform.transform(lr_x_utm10,lr_y_utm10,direction='INVERSE')
print(f"{(ul_lon,ul_lat,lr_lon,lr_lat)=}")

# %% [markdown]
# ## Make the map extent bigger than the image
#
# In order show more context, I'll enlarge the map extent by 5 km on each size

# %%
map_ul_x = ul_x_utm10 - 5.e3
map_lr_x = lr_x_utm10 + 5.e3
map_ul_y = ul_y_utm10 + 5.e3
map_lr_y = lr_y_utm10 - 5.e3
map_extent_utm10 = (map_ul_x,map_lr_x,map_lr_y,map_ul_y)
print(f"{map_extent_utm10=}")

# %% [markdown]
# * Repeat the imshow code from {ref}`rasterio_3bands`:

# %%
cartopy_utm10 = cartopy.crs.epsg(crs_10.to_epsg())
fig, ax = plt.subplots(
        1, 1, figsize=(15,15), subplot_kw={"projection": cartopy_utm10}
    )

vmin = 0.0
vmax = 0.4
the_norm = Normalize(vmin=vmin, vmax=vmax, clip=False)
palette = "viridis"
pal = copy.copy(plt.get_cmap(palette))
pal.set_bad("0.75")  # 75% grey for out-of-map cells
pal.set_over("w")  # color cells > vmax red
pal.set_under("k")  # color cells < vmin black
col=ax.imshow(b5_refl, cmap=pal, norm=the_norm, origin="upper",
          extent=image_extent_utm10,transform=cartopy_utm10);
cbar_ax = fig.add_axes([0.85, 0.2, 0.05, 0.6])
cbar = ax.figure.colorbar(col, extend="both", cax=cbar_ax, orientation="vertical")
cbar.set_label("band 5 reflectance")

# %% [markdown]
# ## Now set the map extent
#
# From the {ref}`demo_cartopy_extent` notebook we can repeat the ax.set_extent call to make the map
# larger than the image

# %%
ax.set_extent(map_extent_utm10,crs=cartopy_utm10)
display(fig)

# %% [markdown]
# ## Add features
#
# Use the {ref}`subset_map` code to put features on this axis.  I need the cartopy_latlon
# crs to do this (since p_latlon, the pyproj crs, doesn't work with cartopy)

# %%
gpd_dict = {}
read_files=True
if read_files:
    all_cia = a301_lib.data_share / "openstreetmap/WDBII_shp/f"
    all_cia = list(all_cia.glob("*"))
    for item in all_cia:
        gpd_dict[item.name] = gpd.read_file(item)
        print(f"read {item.name}")

    all_gshhs = a301_lib.data_share / "openstreetmap/GSHHS_shp/f"
    all_gshhs = list(all_gshhs.glob("*"))

    for item in all_gshhs:
        gpd_dict[item.name] = gpd.read_file(item)
        print(f"read {item.name}")
else:
    shape_files = list(small_shapes.glob("*"))
    for item in shape_files:
        key = item.stem
        gpd_dict[key] = gpd.read_file(item)
        print((f"reading saved shapefile {item} with\n"
               f"{len(gpd_dict[key])} rows"))

# %%
def find_features(extent, df):
    """
    given an extent and a dataframe, return a new dataframe
    containing only features that fall within the extent

    Parameters
    ----------

    extent:  list -- geographic extent in lon (deg E)/lat (deg N)
    df:  the geopandas dataframe to slice
    """
    xleft, xright, ybot, ytop = extent
    hit_rows = df.cx[xleft:xright, ybot:ytop]
    return hit_rows

extent = [-124, -122, 48, 50]
if read_files:
    subset_dict = {}
    for key, df in gpd_dict.items():
        df_subset = find_features(extent, df)
        if len(df_subset) > 0:
            subset_dict[key] = df_subset
            print(f"clipping {key}")
else:
    subset_dict=gpd_dict

# %% [markdown]
# * now put the features on -- they are defined in geodetic (lat/lon) crs

# %%
cartopy_latlon = cartopy.crs.PlateCarree()
for key, df in subset_dict.items():
    print(f"adding {key} with {len(df)} features")
    if key.find("river") > -1:
        ax.add_geometries(
            df["geometry"], cartopy_latlon, facecolor="none", edgecolor="green",lw=3,
        )
    else:
        ax.add_geometries(
            df["geometry"], cartopy_latlon, facecolor="none", edgecolor="blue",lw=3,
        )
display(fig)

# %% [markdown]
# ## Putting images from different utm zones on a map
#
# (not part of the assignment)
#
# Suppose you're working on a project that spans BC and Alberta, and you need
# to put images on a map that are either UTM zone 10 (-126 to -120 deg E) or
# UTM zone 11 (-120 to - 114 deg E).  To show both images on the same map, you
# need to reproject them onto a common crs. As a compromise, let's use a laea
# projection centered on `lon_0`= -120 deg E, `lat_0`=50 deg N
#
# Here are the rasterio reprojection module docs: [rasterio](https://rasterio.readthedocs.io/en/latest/topics/reproject.html)
#
# ### Step 1 is to create [the new pyproj crs](https://pyproj4.github.io/pyproj/dev/api/crs/crs.html).
#
# I'll just borrow the pyproj parameters from our earlier notebooks.

# %%
lat_0=50
lon_0 = -120
laea_proj = {'datum': 'WGS84', 'lat_0': '50', 'lon_0': '-120', 'no_defs': 'None',
             'proj': 'laea', 'type': 'crs', 'units': 'm', 'x_0': '0', 'y_0': '0'}
p_laea = CRS(laea_proj)
p_laea.to_wkt()

# %% [markdown]
# ### Step 2 is to reproject the extent in the new coordinate system
#
# Remember that extent order is:  [xleft, xright, ybot, ytop].  We need to get these values in the new laea crs.
#
# Here are the two extents for the image"

# %%
crs_transform = Transformer.from_crs(p_utm10, p_laea)
extent_utm10 = [ul_x_utm10,lr_x_utm10,lr_y_utm10,ul_y_utm10]
ul_x_laea,ul_y_laea = crs_transform.transform(ul_x_utm10,ul_y_utm10)
lr_x_laea,lr_y_laea = crs_transform.transform(lr_x_utm10,lr_y_utm10)
extent_laea=[ul_x_laea,lr_x_laea,lr_y_laea,ul_y_laea]
print(f"{extent_utm10=}")
print(f"{extent_laea=}")

# %% [markdown]
# ### step 3: create the affine transform
#
# We need the pixel size, then we can use the Affine constructor as we did in {ref}`image_zoom`
# (don't forget to make the pixel height negative). Notice that the pixel size has
# changed slightly in the laea crs.  Also notice that the correct denominator is (#pixels - 1), not #pixels.  That's
# because the extents go from the left side of the first pixel to the right side of the last pixel.  Dividing
# by #pixels would put the right boundary 1 pixel less than it needs to be.

# %%
pixel_x_size = (lr_x_laea - ul_x_laea)/(profile['width'] -1)
pixel_y_size = (ul_y_laea - lr_y_laea)/(profile['height'] - 1)
print(f"{pixel_x_size=},{pixel_y_size=}")

# %%
from affine import Affine

laea_affine = Affine(pixel_x_size,0,ul_x_laea,0,-pixel_y_size,ul_y_laea)

# %% [markdown]
# ### step 4: do the reprojection from p_utm10 to p_laea
#
# First make a numpy array to hold the reprojected image.  We're keeping the row and column numbers the same as i the original tiff.

# %%
width=profile['width']
height=profile['height']
b5_refl_laea = np.ones([height,width],dtype=np.float32)

# %% [markdown]
# ### Now reproject from utm10 to laea

# %%
from rasterio.warp import Resampling, reproject

reproject(
        b5_refl,
        b5_refl_laea,
        src_transform=affine_transform,
        src_crs=p_utm10,
        dst_transform=laea_affine,
        dst_crs=p_laea,
        resampling=Resampling.nearest);
plt.imshow(b5_refl_laea);

# %% [markdown]
# ### step 5: Make a cartopy map
#
# As of cartopy 0.18 we also need to create a cartopy version of the laea crs,
# since it doesn't accept the pyproj version.   See [cartopy projections](https://scitools.org.uk/cartopy/docs/latest/crs/projections.html#cartopy-projections).  One difference is
# that cartopy requires that the datum be specified separately from the projection, using a `globe` object.
# You can track progress on making cartopy more compatible with pyproj [here](https://github.com/SciTools/cartopy/pull/1023#discussion_r168395702) [and here](https://github.com/SciTools/cartopy/issues/1477)

# %%
dir(cartopy.crs.CRS)
globe = cartopy.crs.Globe(datum='WGS84',ellipse='WGS84')
laea_cartopy_crs = cartopy.crs.LambertAzimuthalEqualArea(central_longitude=lon_0,
                                                    central_latitude=lat_0,
                                                    globe=globe)
print(f"proj4 string: {laea_cartopy_crs.proj4_init=}")

# %% [markdown]
# **now plot it with features**
#
# Note the boundaries are slightly skewed at the top and bottom of the reprojection. In the laea projection all locations are referenced to the point at (`lon_`,`lat_0`), while for the UTM projection all locations are referenced to the longitude line going through the middle of the zone. The UTM preserves angles and shapes over small regions (i.e. it's conformal), but the scale changes with location.  The laea preserves scale, but distorts angles and shapes.  See [wikpedia lambert](https://en.wikipedia.org/wiki/Lambert_azimuthal_equal-area_projection) and [wikipedia conformal](https://en.wikipedia.org/wiki/Conformal_map_projection)

# %%
fig, ax = plt.subplots(
        1, 1, figsize=(10,15), subplot_kw={"projection": laea_cartopy_crs}
    )

vmin = 0.0
vmax = 0.4
the_norm = Normalize(vmin=vmin, vmax=vmax, clip=False)
palette = "viridis"
pal = copy.copy(plt.get_cmap(palette))
pal.set_bad("0.75")  # 75% grey for out-of-map cells
pal.set_over("w")  # color cells > vmax red
pal.set_under("k")  # color cells < vmin black
col=ax.imshow(b5_refl_laea, cmap=pal, norm=the_norm, origin="upper",
          extent=extent_laea,transform=laea_cartopy_crs)
cbar_ax = fig.add_axes([0.90, 0.2, 0.05, 0.6])
cbar = ax.figure.colorbar(col, extend="both", cax=cbar_ax, orientation="vertical")
cbar.set_label("band 5 reflectance")

# %% [markdown]
# ### Add a coastline/rivers

# %%
for key, df in subset_dict.items():
    print(f"adding {key} with {len(df)} features")
    if key.find("river") > -1:
        ax.add_geometries(
            df["geometry"], cartopy_latlon, facecolor="none", edgecolor="red",lw=5,
        )
    else:
        ax.add_geometries(
            df["geometry"], cartopy_latlon, facecolor="none", edgecolor="blue",lw=3,
        )
display(fig)

# %% [markdown]
# ## draw a red box in the middle of the scene
#
# We want a box that's 50 rows by 100 columns. I'll center it at row 300, column 200, and move up 25 rows and
# down 50 rows to find the corners.

# %%
ul_x, ul_y = laea_affine*(150,275)
lr_x, lr_y = laea_affine*(250,325)
delta_x = lr_x - ul_x
delta_y = ul_y - lr_y
print(f"{(ul_x,ul_y,lr_x,lr_y)=}")
#
# circle clockwise from upper left corner
#
box_x = [ul_x, ul_x + delta_x, ul_x + delta_x, ul_x,          ul_x]
box_y = [ul_y, ul_y          , ul_y - delta_y, ul_y - delta_y,ul_y]
ax.plot(box_x,box_y,'r-',lw=4)
display(fig)

# %% [markdown]
# ## More mapping tools for meteorology
#
# Check out [metpy](https://unidata.github.io/MetPy/latest/examples/index.html#plotting) for meteorologically oriented maps.


#!/usr/bin/python

'''
NUI Galway CT5132/CT5148 Programming and Tools for AI (James McDermott)

Solution for Assignment 3:  Hand-coding solutions for the Abstraction
                            and Reasoning Corpus

By writing my name below and submitting this file, I/we declare that
all additions to the provided skeleton file are my/our own work, and
that I/we have not seen any work on this assignment by another
student/group.

Student name(s):    John Doyle
Student ID(s):      21230978
GitHub:             https://github.com/john9d/ARC

Summary of the ARC Assignment:

The solve_* solutions completed below were completed using basic python functions
and no machine learning techniques, with the majority using only the built in Python
functions. Such as for loops, if else statements and enumerate. The math library was
used for an irrational number and the numpy library for mainly array specific functions.

Most of the solve_* solution started with a loop pass through the array and getting the
index of specific number required to solve the task. Then mainly replacing this with the
correct number, some requiring to reshape. Where possible the solve_* solutions tried to
use techniques that would help solve the solution in different size array that the tasks
shown.

Below the term "colour" is used primarily, as the ARC testing_interface.html was used
to identify/describe the task and its solution.

'''

import os, sys
import json
import numpy as np
import re
import math

### YOUR CODE HERE: write at least three functions which solve
### specific tasks by transforming the input x and returning the
### result. Name them according to the task ID as in the three
### examples below. Delete the three examples. The tasks you choose
### must be in the data/training directory, not data/evaluation.

def solve_22eb0ac0(x):
    '''
    ########################### SOLUTION No.1 ###############################
    #Using the ARC testing interface Task 22eb0ac0 shows a 10x10 array where there
    #are a number of colours vertically aligned and opposing each other on the
    #left and right side of the array. The task solutions show an array where the
    #a row in the array is filled in wth a colour that had a match of colour on the
    #left and right hand side. Any colours that do not match, will leave the row
    #black in between.

    #To solve this task we will need iterate through the array to adn create an if
    #statement to see if the start of the row and the end of the row match. If this
    #is true, we will enumerate the row, identify and replace all of the colours in
    #between with the match left and right colour.
    '''
    # Initiate loop through the array
    for i in x:

        # We will now check for the matching numbers
        if i[0] == i[-1]:

            # We will need to store the number so we can replace the non-matching ones
            value = i[0]

            # We will store the non-matching number
            item = i[1]

            # Enumerate through the row and create an index of the non-matching values
            index = [j for j, y in enumerate(i) if y == item]

            # We will now loop through the index list so we can replace this values
            for a in index:

                # This will now replace the non-matching values
                i[a] = value

    return x

def solve_0d3d703e(x):
    '''
    ########################### SOLUTION No.2 ###############################
    #Using the ARC testing interface Task 0d3d703e shows a 3x3 array of 3 colours,
    #all vertically aligned. Each colour is converted in the solution i.e. green to
    #yellow, blue to grey, etc. and the conversion is mirrored i.e. yellow to green,
    #grey to blue, etc. The test task has an array of light blue, blue and green all
    #of which need to be converted to wine, grey and yellow respectively.

    #To solve this task we will need to convert each colour (number), however there
    #are overlaps in colours. If we attempt to simply convert each number to its
    #solved state, the number will be reconverted back i.e. when updating 1 to 5, next
    #in the array we are required to update 5 to 1. This will return the old number.
    #To avoid this, we are going to divide by an irrational number (pi in this case)
    #so we can avoid double conversion. Post this, we will multiple the array by pi and
    #perform floor division to remove a chance of a floating point error.
    '''

    # We are going search the array and covert 1 to 5 and divide by pi
    x = np.where(x == 1, 5/math.pi, x)

    # We are going search the array and covert 2 to 6 and divide by pi
    x = np.where(x == 2, 6/math.pi, x)

    # We are going search the array and covert 3 to 4 and divide by pi
    x = np.where(x == 3, 4/math.pi, x)

    # We are going search the array and covert 4 to 3 and divide by pi
    x = np.where(x == 4, 3/math.pi, x)

    # We are going search the array and covert 6 to 2 and divide by pi
    x = np.where(x == 6, 2/math.pi, x)

    # We are going search the array and covert 8 to 9 and divide by pi
    x = np.where(x == 8, 9/math.pi, x)

    # We are going search the array and covert 5 to 1 and divide by pi
    x = np.where(x == 5, 1/math.pi, x)

    # We are going search the array and covert 9 to 8 and divide by pi
    x = np.where(x == 9, 8/math.pi, x)

    # We are going now multiply the array by pi and perform a floor division
    # This will return remove the floating point error
    x = (x*math.pi)//1

    return x

def solve_ce4f8723(x):
    '''
    ########################### SOLUTION No.3 ###############################
    #Using the ARC testing interface Task ce4f8723 shows a 9x4 array. Where there
    #are two 4x4 arrays with coloured and black squares, separated by a 1x4 array
    #to divide the two 4x4 arrays. The results of each task demonstrate that when
    #overlay the two 4x4 arrays, the resulting array shows only black squares in
    #the resulting 4x4 array where both have them. Colouring all of the other squares
    #in a different colour (green!)

    #To solve this task we will need to identify the shape of the array, and then we will
    #flatten the array to (1, N), where N is the total number of items in the array. As
    #we need to combine the two 4x4 coloured arrays we will then split the (1,N) array
    #into two, removing the dividing array. We shall then add the two arrays together and
    #then enumerate it to find where there are black colours (0). We then replace the
    #black (non-zero) with the green and reshape to the correct format for the output.
    '''

    # First we need to define the shape of the array being assessed using .shape
    h, w = x.shape

    # Then we need to flatten nDimension out the array to a (1, N) array
    y = x.ravel()

    # We need to split out the top array of colours from 0 to the width times the width
    # index in the flatten array. This will future proof should we any other NxN array
    y1 = y[0:(w*w)]

    # We will split out the bottom array of colours, removing the 1x4 array by starting
    # the index at width times the width plus the width
    y2 = y[(w*w+w):]

    # Add the two arrays together so we can identify the black (zeros)
    y_total = (y1+y2)

    # Set the value for the green output colour
    value = 3

    # Enumerate through the list to identify the zeros in the combined arrays
    index = [j for j, a in enumerate(y_total) if a != 0]

    # We now need to iterate through the list of non-zeros
    for b in index:

        # Replacing the values with the green (3)
        y_total[b] = value

    # Now we will reshape the (1, N) array back to the correct shape
    x = y_total.reshape(w, -w)

    return x

def solve_e26a3af2(x):
    '''
    ########################### SOLUTION No.4 ###############################
    #Using the ARC testing interface Task e26a3af2 shows a varying NxM arrays.
    #Where some of the arrays N = M. In each of the tasks, there is either a
    #vertical or horizontal lines of varying width containing primarily 3 or 4
    #colours (numbers), with a random placement of any colour throughout. Each
    #task solution clears out the random colours and leaves a neat 3 or 4 colour
    #lines. The test grid is 5 colours in a horizontal arrangement.

    #To solve the task we will enumerate the array and identify the mode colour
    #in each row of the array, inorder to create a new NxM array in clean of any
    #random colours. This will work with the horizontally aligned colours easily,
    #however we will need to transpose the array for the vertically aligned colours.
    #An if statement that will assess if the mode is greater than 1/2 of the array
    #width (as all arrays have 3 or more colours) will be able to ensure the mode
    #and array will be correctly reformed.
    '''

    # We need to create an empty list for mode colour in each row
    mode = []

    # We need to create an empty list for maximum number of colours in each row
    # This will allow us to verify if we need to transpose the array
    total = []

    # Gather the height and width of the array so we can create the new array correctly
    h, w = x.shape

    # We will enumerate x to identify the the mode colour in each row
    for i, j in enumerate(x):

        # Create a list of each value and count of each to identify the mode
        vals, counts = np.unique(x[i], return_counts=True)

        # Capture the value tha made the colour the mode
        total = max(counts)

        # Create an location index for the value which is the mode
        index = np.argmax(counts)

        # Store the mode in the row
        c = vals[index]

        # Append the mode to the global list
        mode = np.append(mode, [c], axis=0)

    # We will now transpose the array and complete the same tasks as above
    # We need to create an empty list for mode colour in each row
    x_T = x.T

    # We need to create an empty list for maximum number of colours in each row
    # This will allow us to verify if we need to transpose the array
    mode_T = []

    # Gather the height and width of the array so we can create the new array correctly
    h_T, w_T = x_T.shape

    # We will enumerate x to identify the the mode colour in each row
    for i_T, j_T in enumerate(x_T):

        # Create a list of each value and count of each to identify the mode
        vals_T, counts_T = np.unique(x_T[i_T], return_counts=True)

        # Create an location index for the value which is the mode
        index_T = np.argmax(counts_T)

        # Store the mode in the row
        c_T = vals_T[index_T]

        # Append the mode to the global list
        mode_T = np.append(mode_T, [c_T], axis=0)

    # Create an if statement to verify if we are working with a vertical task
    if total < w/2:
        ##The mode is less than half the count so we will use the transposed array##
        # Create an empty numpy array
        f_T = np.array([])

        # Iterate through the global mode list created
        for i in mode_T:
            # Append each mode value as a new row and multiply it by the width of
            # of the original transposed array
            f_T = np.append(f_T, [i] * w_T, axis=0)

        # Reshape the new array to the original transposed array
        r = f_T.reshape(h_T, w_T)

        # Transpose the new array back to its original shape
        x = r.T

    # Else we are now working with a horizontal type task
    else:
        ##The mode is more than half the count so we will use the normal array##
        # Create an empty numpy array
        f = np.array([])

        # Iterate through the global mode list created
        for i in mode:

            # Append each mode value as a new row and multiply it by the width of
            # of the original array
            f = np.append(f, [i] * w, axis=0)

        # Reshape the new array to the original array
        x = f.reshape(h, w)

    return x

def solve_9ecd008a(x):
    '''
    ########################### SOLUTION No.5 ###############################
    Using the ARC testing interface Task 9ecd008 shows a large 20x20 array.
    All with a pattern of different colours. Each task has a small 3x3 array with
    colour missing, the solution for each task is to create a 3x3 of the missing
    colours in the correct pattern. As each pattern is mirrored on x and y axes where
    10x20 = 10x20 and 20x10 = 20x10, the missing pattern is its mirror on the y axis
    in each task.

    To solve the task we will enumerate the array and identify which rows contain
    a black (zero) square by calculating the length of the row with no black. If
    the row has black square, we will pass this into another loop to enumerate for
    the index of the black square in the row. In turn getting the absolute difference
    between the index and the length of the row so we capture the index of the mirror
    colours. We will then append the correct colour into a new list and finally
    reshape it to the 3x3 array.
    '''
    # 1st we will create an empty list to append the correct pattern into
    answer = []

    # We will now get the shape of the array to determine the length of each row
    h, w = x.shape

    # We will now start the loop through the array
    for i in x:

        # We will set the value for the black square
        black = 0

        # We will enumerate to identify the length of each row that has a black square
        non_count = [j for j, y in enumerate(i) if y != black]

        # Now create an if statement to pass the row into the next loop if contains a
        # black square
        if len(non_count) != w:

            # We will now enumerate to identify the count of the black squares
            zero_index = [j for j, y in enumerate(i) if y == black]

            # Create a (1,3) array of the maximum index number so we can calculate
            # the mirror location of the black square
            mirror_index = [w - 1] * max(zero_index)

            # Create an empty list of the mirrored index number
            index = []

            # Create a loop to get the location of the black square and the mirror
            # location by getting the absolute difference between the indexes
            for q, p in zip(mirror_index, zero_index):
                index.append(abs(q - p))

            # Finally we will append the colour of the mirrored index in the row
            for c in index:

                # Append the colour to the def global answer
                answer = np.append(answer, [i[c]], axis=0)

    # We will now reshape the list to the correct array size
    x = answer.reshape(3, -3)

    return x

def solve_178fcbfb(x):
    '''
    ########################### SOLUTION No.6 ###############################
    Using the ARC testing interface Task 178fcbfb shows a number of NxM arrays
    where N is sometimes equal to M. There are 3 colours red, green and
    blue in a number of squares in each array. The solution for each task shows
    that the green and blue squares fill their row completely with their colour.
    Then the the red is filled in vertically, under-lapping the green and blue
    squares.

    To solve the task we will enumerate the array and identify which rows contain
    a black (zero) square by calculating the length of the row with a non-black. If
    the row a non-black square, we will pass this into another loop to enumerate for
    the value of the colour in the row. If not red (2) we will then enumerate the row
    to replace all the black (0) squares in with its correct colour. After this we
    will redo the above by transposing the array and look for the red (2) squares
    only to replace the black (0) values with red.
    '''
    # First we need to get the shape of the array to sense check which rows need to
    # be passed into the loop to fill out the colour in the row
    h, w = x.shape

    # Initiate loop through the array
    for i in x:

        # We will set the value for the black square
        black = 0

        # We will enumerate to identify the length of each row that has a black square
        is_black = [j for j, y in enumerate(i) if y == black]

        # Now create an if statement to pass the row into the next loop if contains a
        # black square
        if len(is_black) != w:

            # We will now enumerate to identify the count of the black squares
            non_black = [j for j, y in enumerate(i) if y != black]

            # We will now store the colour value
            colour = i[non_black]

            # We will now sense check that this row does not contain red (2)
            if colour != 2:

                # We will now enumerate to identify the index of the black squares
                zero_index = [j for j, y in enumerate(i) if y == black]

                # We will now loop through to replace the black squares
                for b in zero_index:

                    # Replacing the values with the colour
                    i[b] = colour

    # Initiate loop through the transposed array
    for i in x.T:

        # We will set the value for the black square
        black = 0

        # We will enumerate to identify the length of each row that has a black square
        is_black = [j for j, y in enumerate(i) if y == black]

        # Now create an if statement to pass the row into the next loop if contains a
        # black square
        if len(is_black) != w:

            # We will now enumerate to identify the count of the black squares
            non_black = [j for j, y in enumerate(i) if y != black]

            # We will now store the colour value
            colour = i[non_black]

            # We will need to only pass in the row with the red (2) value
            if 2 in colour:

                # We will now enumerate to identify the count of the black squares
                zero_index = [j for j, y in enumerate(i) if y == black]

                # We will now loop through to replace the black squares
                for b in zero_index:

                    # Replacing the values with the red (2)
                    i[b] = 2
    return x

def main():
    # Find all the functions defined in this file whose names are
    # like solve_abcd1234(), and run them.

    # regex to match solve_* functions and extract task IDs
    p = r"solve_([a-f0-9]{8})"
    tasks_solvers = []
    # globals() gives a dict containing all global names (variables
    # and functions), as name: value pairs.
    for name in globals():
        m = re.match(p, name)
        if m:
            # if the name fits the pattern eg solve_abcd1234
            ID = m.group(1) # just the task ID
            solve_fn = globals()[name] # the fn itself
            tasks_solvers.append((ID, solve_fn))

    for ID, solve_fn in tasks_solvers:
        # for each task, read the data and call test()
        directory = os.path.join("..", "data", "training")
        json_filename = os.path.join(directory, ID + ".json")
        data = read_ARC_JSON(json_filename)
        test(ID, solve_fn, data)

def read_ARC_JSON(filepath):
    """Given a filepath, read in the ARC task data which is in JSON
    format. Extract the train/test input/output pairs of
    grids. Convert each grid to np.array and return train_input,
    train_output, test_input, test_output."""

    # Open the JSON file and load it
    data = json.load(open(filepath))

    # Extract the train/test input/output grids. Each grid will be a
    # list of lists of ints. We convert to Numpy.
    train_input = [np.array(data['train'][i]['input']) for i in range(len(data['train']))]
    train_output = [np.array(data['train'][i]['output']) for i in range(len(data['train']))]
    test_input = [np.array(data['test'][i]['input']) for i in range(len(data['test']))]
    test_output = [np.array(data['test'][i]['output']) for i in range(len(data['test']))]

    return (train_input, train_output, test_input, test_output)

def test(taskID, solve, data):
    """Given a task ID, call the given solve() function on every
    example in the task data."""
    print(taskID)
    train_input, train_output, test_input, test_output = data
    print("Training grids")
    for x, y in zip(train_input, train_output):
        yhat = solve(x)
        show_result(x, y, yhat)
    print("Test grids")
    for x, y in zip(test_input, test_output):
        yhat = solve(x)
        show_result(x, y, yhat)

def show_result(x, y, yhat):
    print("Input")
    print(x)
    print("Correct output")
    print(y)
    print("Our output")
    print(yhat)
    print("Correct?")
    if y.shape != yhat.shape:
        print(f"False. Incorrect shape: {y.shape} v {yhat.shape}")
    else:
        print(np.all(y == yhat))

if __name__ == "__main__": main()


"""Module for compiling codegen output, and wrap the binary for use in
python.

.. note:: To use the autowrap module it must first be imported

   >>> from sympy.utilities.autowrap import autowrap

This module provides a common interface for different external backends, such
as f2py, fwrap, Cython, SWIG(?) etc. (Currently only f2py and Cython are
implemented) The goal is to provide access to compiled binaries of acceptable
performance with a one-button user interface, i.e.

    >>> from sympy.abc import x,y
    >>> expr = ((x - y)**(25)).expand()
    >>> binary_callable = autowrap(expr)
    >>> binary_callable(1, 2)
    -1.0

The callable returned from autowrap() is a binary python function, not a
SymPy object.  If it is desired to use the compiled function in symbolic
expressions, it is better to use binary_function() which returns a SymPy
Function object.  The binary callable is attached as the _imp_ attribute and
invoked when a numerical evaluation is requested with evalf(), or with
lambdify().

    >>> from sympy.utilities.autowrap import binary_function
    >>> f = binary_function('f', expr)
    >>> 2*f(x, y) + y
    y + 2*f(x, y)
    >>> (2*f(x, y) + y).evalf(2, subs={x: 1, y:2})
    0.e-110

The idea is that a SymPy user will primarily be interested in working with
mathematical expressions, and should not have to learn details about wrapping
tools in order to evaluate expressions numerically, even if they are
computationally expensive.

When is this useful?

    1) For computations on large arrays, Python iterations may be too slow,
       and depending on the mathematical expression, it may be difficult to
       exploit the advanced index operations provided by NumPy.

    2) For *really* long expressions that will be called repeatedly, the
       compiled binary should be significantly faster than SymPy's .evalf()

    3) If you are generating code with the codegen utility in order to use
       it in another project, the automatic python wrappers let you test the
       binaries immediately from within SymPy.

    4) To create customized ufuncs for use with numpy arrays.
       See *ufuncify*.

When is this module NOT the best approach?

    1) If you are really concerned about speed or memory optimizations,
       you will probably get better results by working directly with the
       wrapper tools and the low level code.  However, the files generated
       by this utility may provide a useful starting point and reference
       code. Temporary files will be left intact if you supply the keyword
       tempdir="path/to/files/".

    2) If the array computation can be handled easily by numpy, and you
       don't need the binaries for another project.

"""

from __future__ import print_function, division

import sys
import os
import shutil
import tempfile
from subprocess import STDOUT, CalledProcessError, check_output
from string import Template
from warnings import warn

from sympy.core.cache import cacheit
from sympy.core.compatibility import range, iterable
from sympy.core.function import Lambda
from sympy.core.relational import Eq
from sympy.core.symbol import Dummy, Symbol
from sympy.tensor.indexed import Idx, IndexedBase
from sympy.utilities.codegen import (make_routine, get_code_generator,
                                     OutputArgument, InOutArgument,
                                     InputArgument, CodeGenArgumentListError,
                                     Result, ResultBase, C99CodeGen)
from sympy.utilities.lambdify import implemented_function
from sympy.utilities.decorator import doctest_depends_on

_doctest_depends_on = {'exe': ('f2py', 'gfortran', 'gcc'),
                       'modules': ('numpy',)}

class CodeWrapError(Exception):
    pass

class CodeWrapper(object):
    """Base Class for code wrappers"""
    _filename = "wrapped_code"
    _module_basename = "wrapper_module"
    _module_counter = 0

    @property
    def filename(self):
        return "%s_%s" % (self._filename, CodeWrapper._module_counter)

    @property
    def module_name(self):
        return "%s_%s" % (self._module_basename, CodeWrapper._module_counter)

    def __init__(self, generator, filepath=None, flags=[], verbose=False):
        """
        generator -- the code generator to use
        """
        self.generator = generator
        self.filepath = filepath
        self.flags = flags
        self.quiet = not verbose

    @property
    def include_header(self):
        return bool(self.filepath)

    @property
    def include_empty(self):
        return bool(self.filepath)

    def _generate_code(self, main_routine, routines):
        routines.append(main_routine)
        self.generator.write(
            routines, self.filename, True, self.include_header,
            self.include_empty)

    def wrap_code(self, routine, helpers=[]):

        workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
        if not os.access(workdir, os.F_OK):
            os.mkdir(workdir)
        oldwork = os.getcwd()
        os.chdir(workdir)
        try:
            sys.path.append(workdir)
            self._generate_code(routine, helpers)
            self._prepare_files(routine)
            self._process_files(routine)
            mod = __import__(self.module_name)
        finally:
            sys.path.remove(workdir)
            CodeWrapper._module_counter += 1
            os.chdir(oldwork)
            if not self.filepath:
                try:
                    shutil.rmtree(workdir)
                except OSError:
                    # Could be some issues on Windows
                    pass

        return self._get_wrapped_function(mod, routine.name)

    def _process_files(self, routine):
        command = self.command
        command.extend(self.flags)
        try:
            retoutput = check_output(command, stderr=STDOUT)
        except CalledProcessError as e:
            raise CodeWrapError(
                "Error while executing command: %s. Command output is:\n%s" % (
                    " ".join(command), e.output.decode()))
        if not self.quiet:
            print(retoutput)

class DummyWrapper(CodeWrapper):
    """Class used for testing independent of backends """

    template = """# dummy module for testing of SymPy
def %(name)s():
    return "%(expr)s"
%(name)s.args = "%(args)s"
%(name)s.returns = "%(retvals)s"
"""

    def _prepare_files(self, routine):
        return

    def _generate_code(self, routine, helpers):
        with open('%s.py' % self.module_name, 'w') as f:
            printed = ", ".join(
                [str(res.expr) for res in routine.result_variables])
            # convert OutputArguments to return value like f2py
            args = filter(lambda x: not isinstance(
                x, OutputArgument), routine.arguments)
            retvals = []
            for val in routine.result_variables:
                if isinstance(val, Result):
                    retvals.append('nameless')
                else:
                    retvals.append(val.result_var)

            print(DummyWrapper.template % {
                'name': routine.name,
                'expr': printed,
                'args': ", ".join([str(a.name) for a in args]),
                'retvals': ", ".join([str(val) for val in retvals])
            }, end="", file=f)

    def _process_files(self, routine):
        return

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

class CythonCodeWrapper(CodeWrapper):
    """Wrapper that uses Cython"""

    setup_template = """\
try:
    from setuptools import setup
    from setuptools import Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension
from Cython.Build import cythonize
cy_opts = {cythonize_options}
{np_import}
ext_mods = [Extension(
    {ext_args},
    include_dirs={include_dirs},
    library_dirs={library_dirs},
    libraries={libraries},
    extra_compile_args={extra_compile_args},
    extra_link_args={extra_link_args}
)]
setup(ext_modules=cythonize(ext_mods, **cy_opts))
"""

    pyx_imports = (
        "import numpy as np\n"
        "cimport numpy as np\n\n")

    pyx_header = (
        "cdef extern from '{header_file}.h':\n"
        "    {prototype}\n\n")

    pyx_func = (
        "def {name}_c({arg_string}):\n"
        "\n"
        "{declarations}"
        "{body}")

    std_compile_flag = '-std=c99'

    def __init__(self, *args, **kwargs):
        """Instantiates a Cython code wrapper.

        The following optional parameters get passed to ``distutils.Extension``
        for building the Python extension module. Read its documentation to
        learn more.

        Parameters
        ==========
        include_dirs : [list of strings]
            A list of directories to search for C/C++ header files (in Unix
            form for portability).
        library_dirs : [list of strings]
            A list of directories to search for C/C++ libraries at link time.
        libraries : [list of strings]
            A list of library names (not filenames or paths) to link against.
        extra_compile_args : [list of strings]
            Any extra platform- and compiler-specific information to use when
            compiling the source files in 'sources'.  For platforms and
            compilers where "command line" makes sense, this is typically a
            list of command-line arguments, but for other platforms it could be
            anything. Note that the attribute ``std_compile_flag`` will be
            appended to this list.
        extra_link_args : [list of strings]
            Any extra platform- and compiler-specific information to use when
            linking object files together to create the extension (or to create
            a new static Python interpreter). Similar interpretation as for
            'extra_compile_args'.
        cythonize_options : [dictionary]
            Keyword arguments passed on to cythonize.

        """

        self._include_dirs = kwargs.pop('include_dirs', [])
        self._library_dirs = kwargs.pop('library_dirs', [])
        self._libraries = kwargs.pop('libraries', [])
        self._extra_compile_args = kwargs.pop('extra_compile_args', [])
        self._extra_compile_args.append(self.std_compile_flag)
        self._extra_link_args = kwargs.pop('extra_link_args', [])
        self._cythonize_options = kwargs.pop('cythonize_options', {})

        self._need_numpy = False

        super(CythonCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        command = [sys.executable, "setup.py", "build_ext", "--inplace"]
        return command

    def _prepare_files(self, routine, build_dir=os.curdir):
        # NOTE : build_dir is used for testing purposes.
        pyxfilename = self.module_name + '.pyx'
        codefilename = "%s.%s" % (self.filename, self.generator.code_extension)

        # pyx
        with open(pyxfilename, 'w') as f:
            self.dump_pyx([routine], f, self.filename)

        # setup.py
        ext_args = [repr(self.module_name), repr([pyxfilename, codefilename])]
        if self._need_numpy:
            np_import = 'import numpy as np\n'
            self._include_dirs.append('np.get_include()')
        else:
            np_import = ''

        with open(os.path.join(build_dir, 'setup.py'), 'w') as f:
            includes = str(self._include_dirs).replace("'np.get_include()'",
                                                       'np.get_include()')
            f.write(self.setup_template.format(
                ext_args=", ".join(ext_args),
                np_import=np_import,
                include_dirs=includes,
                library_dirs=self._library_dirs,
                libraries=self._libraries,
                extra_compile_args=self._extra_compile_args,
                extra_link_args=self._extra_link_args,
                cythonize_options=self._cythonize_options
            ))

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name + '_c')

    def dump_pyx(self, routines, f, prefix):
        """Write a Cython file with python wrappers

        This file contains all the definitions of the routines in c code and
        refers to the header file.

        Arguments
        ---------
        routines
            List of Routine instances
        f
            File-like object to write the file to
        prefix
            The filename prefix, used to refer to the proper header file.
            Only the basename of the prefix is used.
        """
        headers = []
        functions = []
        for routine in routines:
            prototype = self.generator.get_prototype(routine)

            # C Function Header Import
            headers.append(self.pyx_header.format(header_file=prefix,
                                                  prototype=prototype))

            # Partition the C function arguments into categories
            py_rets, py_args, py_loc, py_inf = self._partition_args(routine.arguments)

            # Function prototype
            name = routine.name
            arg_string = ", ".join(self._prototype_arg(arg) for arg in py_args)

            # Local Declarations
            local_decs = []
            for arg, val in py_inf.items():
                proto = self._prototype_arg(arg)
                mat, ind = val
                local_decs.append("    cdef {0} = {1}.shape[{2}]".format(proto, mat, ind))
            local_decs.extend(["    cdef {0}".format(self._declare_arg(a)) for a in py_loc])
            declarations = "\n".join(local_decs)
            if declarations:
                declarations = declarations + "\n"

            # Function Body
            args_c = ", ".join([self._call_arg(a) for a in routine.arguments])
            rets = ", ".join([str(r.name) for r in py_rets])
            if routine.results:
                body = '    return %s(%s)' % (routine.name, args_c)
                if rets:
                    body = body + ', ' + rets
            else:
                body = '    %s(%s)\n' % (routine.name, args_c)
                body = body + '    return ' + rets

            functions.append(self.pyx_func.format(name=name, arg_string=arg_string,
                    declarations=declarations, body=body))

        # Write text to file
        if self._need_numpy:
            # Only import numpy if required
            f.write(self.pyx_imports)
        f.write('\n'.join(headers))
        f.write('\n'.join(functions))

    def _partition_args(self, args):
        """Group function arguments into categories."""
        py_args = []
        py_returns = []
        py_locals = []
        py_inferred = {}
        for arg in args:
            if isinstance(arg, OutputArgument):
                py_returns.append(arg)
                py_locals.append(arg)
            elif isinstance(arg, InOutArgument):
                py_returns.append(arg)
                py_args.append(arg)
            else:
                py_args.append(arg)
        # Find arguments that are array dimensions. These can be inferred
        # locally in the Cython code.
            if isinstance(arg, (InputArgument, InOutArgument)) and arg.dimensions:
                dims = [d[1] + 1 for d in arg.dimensions]
                sym_dims = [(i, d) for (i, d) in enumerate(dims) if
                            isinstance(d, Symbol)]
                for (i, d) in sym_dims:
                    py_inferred[d] = (arg.name, i)
        for arg in args:
            if arg.name in py_inferred:
                py_inferred[arg] = py_inferred.pop(arg.name)
        # Filter inferred arguments from py_args
        py_args = [a for a in py_args if a not in py_inferred]
        return py_returns, py_args, py_locals, py_inferred

    def _prototype_arg(self, arg):
        mat_dec = "np.ndarray[{mtype}, ndim={ndim}] {name}"
        np_types = {'double': 'np.double_t',
                    'int': 'np.int_t'}
        t = arg.get_datatype('c')
        if arg.dimensions:
            self._need_numpy = True
            ndim = len(arg.dimensions)
            mtype = np_types[t]
            return mat_dec.format(mtype=mtype, ndim=ndim, name=arg.name)
        else:
            return "%s %s" % (t, str(arg.name))

    def _declare_arg(self, arg):
        proto = self._prototype_arg(arg)
        if arg.dimensions:
            shape = '(' + ','.join(str(i[1] + 1) for i in arg.dimensions) + ')'
            return proto + " = np.empty({shape})".format(shape=shape)
        else:
            return proto + " = 0"

    def _call_arg(self, arg):
        if arg.dimensions:
            t = arg.get_datatype('c')
            return "<{0}*> {1}.data".format(t, arg.name)
        elif isinstance(arg, ResultBase):
            return "&{0}".format(arg.name)
        else:
            return str(arg.name)

class F2PyCodeWrapper(CodeWrapper):
    """Wrapper that uses f2py"""

    def __init__(self, *args, **kwargs):

        ext_keys = ['include_dirs', 'library_dirs', 'libraries',
                    'extra_compile_args', 'extra_link_args']
        msg = ('The compilation option kwarg {} is not supported with the f2py '
               'backend.')

        for k in ext_keys:
            if k in kwargs.keys():
                warn(msg.format(k))
            kwargs.pop(k, None)

        super(F2PyCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        filename = self.filename + '.' + self.generator.code_extension
        args = ['-c', '-m', self.module_name, filename]
        command = [sys.executable, "-c", "import numpy.f2py as f2py2e;f2py2e.main()"]+args
        return command

    def _prepare_files(self, routine):
        pass

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

# Here we define a lookup of backends -> tuples of languages. For now, each
# tuple is of length 1, but if a backend supports more than one language,
# the most preferable language is listed first.
_lang_lookup = {'CYTHON': ('C99', 'C89', 'C'),
                'F2PY': ('F95',),
                'NUMPY': ('C99', 'C89', 'C'),
                'DUMMY': ('F95',)}     # Dummy here just for testing

def _infer_language(backend):
    """For a given backend, return the top choice of language"""
    langs = _lang_lookup.get(backend.upper(), False)
    if not langs:
        raise ValueError("Unrecognized backend: " + backend)
    return langs[0]

def _validate_backend_language(backend, language):
    """Throws error if backend and language are incompatible"""
    langs = _lang_lookup.get(backend.upper(), False)
    if not langs:
        raise ValueError("Unrecognized backend: " + backend)
    if language.upper() not in langs:
        raise ValueError(("Backend {0} and language {1} are "
                          "incompatible").format(backend, language))

@cacheit
@doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
def autowrap(expr, language=None, backend='f2py', tempdir=None, args=None,
             flags=None, verbose=False, helpers=None, code_gen=None, **kwargs):
    """Generates python callable binaries based on the math expression.

    Parameters
    ----------
    expr
        The SymPy expression that should be wrapped as a binary routine.
    language : string, optional
        If supplied, (options: 'C' or 'F95'), specifies the language of the
        generated code. If ``None`` [default], the language is inferred based
        upon the specified backend.
    backend : string, optional
        Backend used to wrap the generated code. Either 'f2py' [default],
        or 'cython'.
    tempdir : string, optional
        Path to directory for temporary files. If this argument is supplied,
        the generated code and the wrapper input files are left intact in the
        specified path.
    args : iterable, optional
        An ordered iterable of symbols. Specifies the argument sequence for the
        function.
    flags : iterable, optional
        Additional option flags that will be passed to the backend.
    verbose : bool, optional
        If True, autowrap will not mute the command line backends. This can be
        helpful for debugging.
    helpers : iterable, optional
        Used to define auxillary expressions needed for the main expr. If the
        main expression needs to call a specialized function it should be put
        in the ``helpers`` iterable. Autowrap will then make sure that the
        compiled main expression can link to the helper routine. Items should
        be tuples with (<funtion_name>, <sympy_expression>, <arguments>). It
        is mandatory to supply an argument sequence to helper routines.
    code_gen : CodeGen instance
        An instance of a CodeGen subclass. Overrides ``language``.
    include_dirs : [string]
        A list of directories to search for C/C++ header files (in Unix form
        for portability).
    library_dirs : [string]
        A list of directories to search for C/C++ libraries at link time.
    libraries : [string]
        A list of library names (not filenames or paths) to link against.
    extra_compile_args : [string]
        Any extra platform- and compiler-specific information to use when
        compiling the source files in 'sources'.  For platforms and compilers
        where "command line" makes sense, this is typically a list of
        command-line arguments, but for other platforms it could be anything.
    extra_link_args : [string]
        Any extra platform- and compiler-specific information to use when
        linking object files together to create the extension (or to create a
        new static Python interpreter).  Similar interpretation as for
        'extra_compile_args'.

    Examples
    --------
    >>> from sympy.abc import x, y, z
    >>> from sympy.utilities.autowrap import autowrap
    >>> expr = ((x - y + z)**(13)).expand()
    >>> binary_func = autowrap(expr)
    >>> binary_func(1, 4, 2)
    -1.0
    """
    if language:
        if not isinstance(language, type):
            _validate_backend_language(backend, language)
    else:
        language = _infer_language(backend)

    helpers = [helpers] if helpers else ()
    args = list(args) if iterable(args, exclude=set) else args

    if code_gen is None:
        code_gen = get_code_generator(language, "autowrap")

    CodeWrapperClass = {
        'F2PY': F2PyCodeWrapper,
        'CYTHON': CythonCodeWrapper,
        'DUMMY': DummyWrapper
    }[backend.upper()]
    code_wrapper = CodeWrapperClass(code_gen, tempdir, flags if flags else (),
                                    verbose, **kwargs)

    helps = []
    for name_h, expr_h, args_h in helpers:
        helps.append(make_routine(name_h, expr_h, args_h))

    for name_h, expr_h, args_h in helpers:
        if expr.has(expr_h):
            name_h = binary_function(name_h, expr_h, backend='dummy')
            expr = expr.subs(expr_h, name_h(*args_h))
    try:
        routine = make_routine('autofunc', expr, args)
    except CodeGenArgumentListError as e:
        # if all missing arguments are for pure output, we simply attach them
        # at the end and try again, because the wrappers will silently convert
        # them to return values anyway.
        new_args = []
        for missing in e.missing_args:
            if not isinstance(missing, OutputArgument):
                raise
            new_args.append(missing.name)
        routine = make_routine('autofunc', expr, args + new_args)

    return code_wrapper.wrap_code(routine, helpers=helps)

@doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
def binary_function(symfunc, expr, **kwargs):
    """Returns a sympy function with expr as binary implementation

    This is a convenience function that automates the steps needed to
    autowrap the SymPy expression and attaching it to a Function object
    with implemented_function().

    Parameters
    ----------
    symfunc : sympy Function
        The function to bind the callable to.
    expr : sympy Expression
        The expression used to generate the function.
    kwargs : dict
        Any kwargs accepted by autowrap.

    Examples
    --------
    >>> from sympy.abc import x, y
    >>> from sympy.utilities.autowrap import binary_function
    >>> expr = ((x - y)**(25)).expand()
    >>> f = binary_function('f', expr)
    >>> type(f)
    <class 'sympy.core.function.UndefinedFunction'>
    >>> 2*f(x, y)
    2*f(x, y)
    >>> f(x, y).evalf(2, subs={x: 1, y: 2})
    -1.0
    """
    binary = autowrap(expr, **kwargs)
    return implemented_function(symfunc, binary)

#################################################################
#                           UFUNCIFY                            #
#################################################################

_ufunc_top = Template("""\
#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/halffloat.h"
#include ${include_file}

static PyMethodDef ${module}Methods[] = {
        {NULL, NULL, 0, NULL}
};""")

_ufunc_outcalls = Template("*((double *)out${outnum}) = ${funcname}(${call_args});")

_ufunc_body = Template("""\
static void ${funcname}_ufunc(char **args, npy_intp *dimensions, npy_intp* steps, void* data)
{
    npy_intp i;
    npy_intp n = dimensions[0];
    ${declare_args}
    ${declare_steps}
    for (i = 0; i < n; i++) {
        ${outcalls}
        ${step_increments}
    }
}
PyUFuncGenericFunction ${funcname}_funcs[1] = {&${funcname}_ufunc};
static char ${funcname}_types[${n_types}] = ${types}
static void *${funcname}_data[1] = {NULL};""")

_ufunc_bottom = Template("""\
#if PY_VERSION_HEX >= 0x03000000
static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "${module}",
    NULL,
    -1,
    ${module}Methods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC PyInit_${module}(void)
{
    PyObject *m, *d;
    ${function_creation}
    m = PyModule_Create(&moduledef);
    if (!m) {
        return NULL;
    }
    import_array();
    import_umath();
    d = PyModule_GetDict(m);
    ${ufunc_init}
    return m;
}
#else
PyMODINIT_FUNC init${module}(void)
{
    PyObject *m, *d;
    ${function_creation}
    m = Py_InitModule("${module}", ${module}Methods);
    if (m == NULL) {
        return;
    }
    import_array();
    import_umath();
    d = PyModule_GetDict(m);
    ${ufunc_init}
}
#endif\
""")

_ufunc_init_form = Template("""\
ufunc${ind} = PyUFunc_FromFuncAndData(${funcname}_funcs, ${funcname}_data, ${funcname}_types, 1, ${n_in}, ${n_out},
            PyUFunc_None, "${module}", ${docstring}, 0);
    PyDict_SetItemString(d, "${funcname}", ufunc${ind});
    Py_DECREF(ufunc${ind});""")

_ufunc_setup = Template("""\
def configuration(parent_package='', top_path=None):
    import numpy
    from numpy.distutils.misc_util import Configuration

    config = Configuration('',
                           parent_package,
                           top_path)
    config.add_extension('${module}', sources=['${module}.c', '${filename}.c'])

    return config

if __name__ == "__main__":
    from numpy.distutils.core import setup
    setup(configuration=configuration)""")

class UfuncifyCodeWrapper(CodeWrapper):
    """Wrapper for Ufuncify"""

    def __init__(self, *args, **kwargs):

        ext_keys = ['include_dirs', 'library_dirs', 'libraries',
                    'extra_compile_args', 'extra_link_args']
        msg = ('The compilation option kwarg {} is not supported with the numpy'
               ' backend.')

        for k in ext_keys:
            if k in kwargs.keys():
                warn(msg.format(k))
            kwargs.pop(k, None)

        super(UfuncifyCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        command = [sys.executable, "setup.py", "build_ext", "--inplace"]
        return command

    def wrap_code(self, routines, helpers=None):
        # This routine overrides CodeWrapper because we can't assume funcname == routines[0].name
        # Therefore we have to break the CodeWrapper private API.
        # There isn't an obvious way to extend multi-expr support to
        # the other autowrap backends, so we limit this change to ufuncify.
        helpers = helpers if helpers is not None else []
        # We just need a consistent name
        funcname = 'wrapped_' + str(id(routines) + id(helpers))

        workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
        if not os.access(workdir, os.F_OK):
            os.mkdir(workdir)
        oldwork = os.getcwd()
        os.chdir(workdir)
        try:
            sys.path.append(workdir)
            self._generate_code(routines, helpers)
            self._prepare_files(routines, funcname)
            self._process_files(routines)
            mod = __import__(self.module_name)
        finally:
            sys.path.remove(workdir)
            CodeWrapper._module_counter += 1
            os.chdir(oldwork)
            if not self.filepath:
                try:
                    shutil.rmtree(workdir)
                except OSError:
                    # Could be some issues on Windows
                    pass

        return self._get_wrapped_function(mod, funcname)

    def _generate_code(self, main_routines, helper_routines):
        all_routines = main_routines + helper_routines
        self.generator.write(
            all_routines, self.filename, True, self.include_header,
            self.include_empty)

    def _prepare_files(self, routines, funcname):

        # C
        codefilename = self.module_name + '.c'
        with open(codefilename, 'w') as f:
            self.dump_c(routines, f, self.filename, funcname=funcname)

        # setup.py
        with open('setup.py', 'w') as f:
            self.dump_setup(f)

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

    def dump_setup(self, f):
        setup = _ufunc_setup.substitute(module=self.module_name,
                                        filename=self.filename)
        f.write(setup)

    def dump_c(self, routines, f, prefix, funcname=None):
        """Write a C file with python wrappers

        This file contains all the definitions of the routines in c code.

        Arguments
        ---------
        routines
            List of Routine instances
        f
            File-like object to write the file to
        prefix
            The filename prefix, used to name the imported module.
        funcname
            Name of the main function to be returned.
        """
        if (funcname is None) and (len(routines) == 1):
            funcname = routines[0].name
        elif funcname is None:
            msg = 'funcname must be specified for multiple output routines'
            raise ValueError(msg)
        functions = []
        function_creation = []
        ufunc_init = []
        module = self.module_name
        include_file = "\"{0}.h\"".format(prefix)
        top = _ufunc_top.substitute(include_file=include_file, module=module)

        name = funcname

        # Partition the C function arguments into categories
        # Here we assume all routines accept the same arguments
        r_index = 0
        py_in, _ = self._partition_args(routines[0].arguments)
        n_in = len(py_in)
        n_out = len(routines)

        # Declare Args
        form = "char *{0}{1} = args[{2}];"
        arg_decs = [form.format('in', i, i) for i in range(n_in)]
        arg_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
        declare_args = '\n    '.join(arg_decs)

        # Declare Steps
        form = "npy_intp {0}{1}_step = steps[{2}];"
        step_decs = [form.format('in', i, i) for i in range(n_in)]
        step_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
        declare_steps = '\n    '.join(step_decs)

        # Call Args
        form = "*(double *)in{0}"
        call_args = ', '.join([form.format(a) for a in range(n_in)])

        # Step Increments
        form = "{0}{1} += {0}{1}_step;"
        step_incs = [form.format('in', i) for i in range(n_in)]
        step_incs.extend([form.format('out', i, i) for i in range(n_out)])
        step_increments = '\n        '.join(step_incs)

        # Types
        n_types = n_in + n_out
        types = "{" + ', '.join(["NPY_DOUBLE"]*n_types) + "};"

        # Docstring
        docstring = '"Created in SymPy with Ufuncify"'

        # Function Creation
        function_creation.append("PyObject *ufunc{0};".format(r_index))

        # Ufunc initialization
        init_form = _ufunc_init_form.substitute(module=module,
                                                funcname=name,
                                                docstring=docstring,
                                                n_in=n_in, n_out=n_out,
                                                ind=r_index)
        ufunc_init.append(init_form)

        outcalls = [_ufunc_outcalls.substitute(
            outnum=i, call_args=call_args, funcname=routines[i].name) for i in
            range(n_out)]

        body = _ufunc_body.substitute(module=module, funcname=name,
                                      declare_args=declare_args,
                                      declare_steps=declare_steps,
                                      call_args=call_args,
                                      step_increments=step_increments,
                                      n_types=n_types, types=types,
                                      outcalls='\n        '.join(outcalls))
        functions.append(body)

        body = '\n\n'.join(functions)
        ufunc_init = '\n    '.join(ufunc_init)
        function_creation = '\n    '.join(function_creation)
        bottom = _ufunc_bottom.substitute(module=module,
                                          ufunc_init=ufunc_init,
                                          function_creation=function_creation)
        text = [top, body, bottom]
        f.write('\n\n'.join(text))

    def _partition_args(self, args):
        """Group function arguments into categories."""
        py_in = []
        py_out = []
        for arg in args:
            if isinstance(arg, OutputArgument):
                py_out.append(arg)
            elif isinstance(arg, InOutArgument):
                raise ValueError("Ufuncify doesn't support InOutArguments")
            else:
                py_in.append(arg)
        return py_in, py_out

@cacheit
@doctest_depends_on(exe=('f2py', 'gfortran', 'gcc'), modules=('numpy',))
def ufuncify(args, expr, language=None, backend='numpy', tempdir=None,
             flags=None, verbose=False, helpers=None, **kwargs):
    """Generates a binary function that supports broadcasting on numpy arrays.

    Parameters
    ----------
    args : iterable
        Either a Symbol or an iterable of symbols. Specifies the argument
        sequence for the function.
    expr
        A SymPy expression that defines the element wise operation.
    language : string, optional
        If supplied, (options: 'C' or 'F95'), specifies the language of the
        generated code. If ``None`` [default], the language is inferred based
        upon the specified backend.
    backend : string, optional
        Backend used to wrap the generated code. Either 'numpy' [default],
        'cython', or 'f2py'.
    tempdir : string, optional
        Path to directory for temporary files. If this argument is supplied,
        the generated code and the wrapper input files are left intact in
        the specified path.
    flags : iterable, optional
        Additional option flags that will be passed to the backend.
    verbose : bool, optional
        If True, autowrap will not mute the command line backends. This can
        be helpful for debugging.
    helpers : iterable, optional
        Used to define auxillary expressions needed for the main expr. If
        the main expression needs to call a specialized function it should
        be put in the ``helpers`` iterable. Autowrap will then make sure
        that the compiled main expression can link to the helper routine.
        Items should be tuples with (<funtion_name>, <sympy_expression>,
        <arguments>). It is mandatory to supply an argument sequence to
        helper routines.
    kwargs : dict
        These kwargs will be passed to autowrap if the `f2py` or `cython`
        backend is used and ignored if the `numpy` backend is used.

    Note
    ----
    The default backend ('numpy') will create actual instances of
    ``numpy.ufunc``. These support ndimensional broadcasting, and implicit type
    conversion. Use of the other backends will result in a "ufunc-like"
    function, which requires equal length 1-dimensional arrays for all
    arguments, and will not perform any type conversions.

    References
    ----------
    [1] http://docs.scipy.org/doc/numpy/reference/ufuncs.html

    Examples
    ========

    >>> from sympy.utilities.autowrap import ufuncify
    >>> from sympy.abc import x, y
    >>> import numpy as np
    >>> f = ufuncify((x, y), y + x**2)
    >>> type(f)
    <class 'numpy.ufunc'>
    >>> f([1, 2, 3], 2)
    array([  3.,   6.,  11.])
    >>> f(np.arange(5), 3)
    array([  3.,   4.,   7.,  12.,  19.])

    For the 'f2py' and 'cython' backends, inputs are required to be equal length
    1-dimensional arrays. The 'f2py' backend will perform type conversion, but
    the Cython backend will error if the inputs are not of the expected type.

    >>> f_fortran = ufuncify((x, y), y + x**2, backend='f2py')
    >>> f_fortran(1, 2)
    array([ 3.])
    >>> f_fortran(np.array([1, 2, 3]), np.array([1.0, 2.0, 3.0]))
    array([  2.,   6.,  12.])
    >>> f_cython = ufuncify((x, y), y + x**2, backend='Cython')
    >>> f_cython(1, 2)  # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    TypeError: Argument '_x' has incorrect type (expected numpy.ndarray, got int)
    >>> f_cython(np.array([1.0]), np.array([2.0]))
    array([ 3.])
    """

    if isinstance(args, Symbol):
        args = (args,)
    else:
        args = tuple(args)

    if language:
        _validate_backend_language(backend, language)
    else:
        language = _infer_language(backend)

    helpers = helpers if helpers else ()
    flags = flags if flags else ()

    if backend.upper() == 'NUMPY':
        # maxargs is set by numpy compile-time constant NPY_MAXARGS
        # If a future version of numpy modifies or removes this restriction
        # this variable should be changed or removed
        maxargs = 32
        helps = []
        for name, expr, args in helpers:
            helps.append(make_routine(name, expr, args))
        code_wrapper = UfuncifyCodeWrapper(C99CodeGen("ufuncify"), tempdir,
                                           flags, verbose)
        if not isinstance(expr, (list, tuple)):
            expr = [expr]
        if len(expr) == 0:
            raise ValueError('Expression iterable has zero length')
        if (len(expr) + len(args)) > maxargs:
            msg = ('Cannot create ufunc with more than {0} total arguments: '
                   'got {1} in, {2} out')
            raise ValueError(msg.format(maxargs, len(args), len(expr)))
        routines = [make_routine('autofunc{}'.format(idx), exprx, args) for
                    idx, exprx in enumerate(expr)]
        return code_wrapper.wrap_code(routines, helpers=helps)
    else:
        # Dummies are used for all added expressions to prevent name clashes
        # within the original expression.
        y = IndexedBase(Dummy())
        m = Dummy(integer=True)
        i = Idx(Dummy(integer=True), m)
        f = implemented_function(Dummy().name, Lambda(args, expr))
        # For each of the args create an indexed version.
        indexed_args = [IndexedBase(Dummy(str(a))) for a in args]
        # Order the arguments (out, args, dim)
        args = [y] + indexed_args + [m]
        args_with_indices = [a[i] for a in indexed_args]
        return autowrap(Eq(y[i], f(*args_with_indices)), language, backend,
                        tempdir, args, flags, verbose, helpers, **kwargs)