"""Native LiteRT lowerings for shape-dependent reductions and shape arithmetic."""
import operator
import torch
from litert_torch import fx_infra
from litert_torch.backend.experimental import torch_tfl
from litert_torch.backend.experimental.torch_tfl._lowerings import _ir_operation
from litert_torch.backend.lowerings import registry,utils
from litert_converter.mlir import ir

fx_infra.decomp.update_pre_lower_decomp(torch_tfl.decomps)

def dynamic_layer_norm(x,normalized_shape,weight,bias,eps):
    axes=list(range(x.ndim-len(normalized_shape),x.ndim))
    mean=x.mean(axes,keepdim=True)
    variance=((x-mean)*(x-mean)).mean(axes,keepdim=True)
    invstd=torch.rsqrt(variance+eps)
    output=(x-mean)*invstd
    if weight is not None:output=output*weight
    if bias is not None:output=output+bias
    return output,mean,invstd

fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.native_layer_norm.default:dynamic_layer_norm})

@registry.lower(operator.floordiv)
def lower_floor_div(lctx,lhs,rhs):
    if isinstance(lhs,int) and isinstance(rhs,int):return lhs//rhs
    if isinstance(lhs,int):lhs=utils.splat(lhs,rhs.type.element_type,rhs.type.shape)
    if isinstance(rhs,int):rhs=utils.splat(rhs,lhs.type.element_type,lhs.type.shape)
    if lhs.type!=rhs.type:raise TypeError('Shape floor division requires matching scalar integer types')
    return _ir_operation('tfl.floor_div',results=[lhs.type],operands=[lhs,rhs])

@registry.lower(torch.ops.aten.relu.default)
def lower_relu(lctx,x):
    return _ir_operation('tfl.relu',results=utils.node_meta_to_ir_types(lctx.node),operands=[x])

@registry.lower(torch.ops.aten.cumsum.default)
def lower_cumsum(lctx,x,dim,dtype=None):
    result=utils.node_meta_to_ir_types(lctx.node)
    if x.type.element_type!=result[0].element_type:
        typ=ir.RankedTensorType.get(x.type.shape,result[0].element_type)
        x=_ir_operation('tfl.cast',results=[typ],operands=[x])
    axis=utils.splat(dim,ir.IntegerType.get_signless(32),[])
    return _ir_operation('tfl.cumsum',results=result,operands=[x,axis],attributes={'exclusive':ir.BoolAttr.get(False),'reverse':ir.BoolAttr.get(False)})

@registry.lower(torch.ops.aten._is_all_true.default)
def lower_all_true(lctx,x):
    if len(x.type.shape)==0:return x
    import numpy as np
    axes=utils.numpy_array_constant(np.arange(len(x.type.shape),dtype=np.int32))
    return _ir_operation('tfl.reduce_all',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,axes],attributes={'keep_dims':ir.BoolAttr.get(False)})

# Preserve symbolic dimensions in the operator schema. The upstream int[]
# schema can specialize reshape dimensions during decomposition.
@torch.library.custom_op('vds::reshape',mutates_args=(),schema='(Tensor x, SymInt[] shape) -> Tensor')
def dynamic_reshape(x,shape):
    return x.reshape(shape).clone()
@dynamic_reshape.register_fake
def dynamic_reshape_fake(x,shape):
    return x.reshape(shape).clone()
@registry.lower(torch.ops.vds.reshape.default)
def lower_reshape(lctx,x,shape):
    return _ir_operation('tfl.reshape',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,utils.convert_shape_to_ir_value(shape)])

def dynamic_view(x,shape):return torch.ops.vds.reshape(x,shape)
fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.view.default:dynamic_view,torch.ops.aten._unsafe_view.default:dynamic_view})

def dynamic_arange(start,end,step=1,*,dtype=None,layout=None,device=None,pin_memory=None):
    result=torch.ops.tfl.range(start,end,step)
    return result.to(dtype=dtype) if dtype is not None else result
fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.arange.start_step:dynamic_arange})

@torch.library.custom_op('vds::broadcast_to',mutates_args=(),schema='(Tensor x, SymInt[] shape) -> Tensor')
def dynamic_broadcast(x,shape):return x.expand(shape).clone()
@dynamic_broadcast.register_fake
def dynamic_broadcast_fake(x,shape):return x.expand(shape).clone()
@registry.lower(torch.ops.vds.broadcast_to.default)
def lower_broadcast(lctx,x,shape):
    return _ir_operation('tfl.broadcast_to',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,utils.convert_shape_to_ir_value(shape)])
fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.expand.default:lambda x,shape,implicit=False:torch.ops.vds.broadcast_to(x,shape)})

@torch.library.custom_op('vds::linear',mutates_args=())
def dynamic_linear(x:torch.Tensor,weight:torch.Tensor,bias:torch.Tensor|None)->torch.Tensor:
    return torch.nn.functional.linear(x,weight,bias)
@dynamic_linear.register_fake
def dynamic_linear_fake(x,weight,bias):
    return x.new_empty(x.shape[:-1] if weight.ndim==1 else (*x.shape[:-1],weight.shape[0]))
@registry.lower(torch.ops.vds.linear.default)
def lower_linear(lctx,x,weight,bias):
    if len(weight.type.shape)==1:
        from litert_converter.mlir.dialects import stablehlo
        from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_squeeze_lowering
        weight=stablehlo.reshape(ir.RankedTensorType.get([1,weight.type.shape[0]],weight.type.element_type),weight)
        if bias is None:bias=utils.splat(0.0,weight.type.element_type,[1])
        typ=utils.node_meta_to_ir_types(lctx.node)[0]
        expanded=ir.RankedTensorType.get([*typ.shape,1],typ.element_type)
        y=_ir_operation('tfl.fully_connected',results=[expanded],operands=[x,weight,bias],attributes={'fused_activation_function':ir.StringAttr.get('NONE'),'weights_format':ir.StringAttr.get('DEFAULT'),'keep_num_dims':ir.BoolAttr.get(True),'asymmetric_quantize_inputs':ir.BoolAttr.get(False)})
        return _tfl_squeeze_lowering(lctx,y,[-1])
    if bias is None:bias=utils.splat(0.0,weight.type.element_type,[weight.type.shape[0]])
    return _ir_operation('tfl.fully_connected',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,weight,bias],attributes={'fused_activation_function':ir.StringAttr.get('NONE'),'weights_format':ir.StringAttr.get('DEFAULT'),'keep_num_dims':ir.BoolAttr.get(True),'asymmetric_quantize_inputs':ir.BoolAttr.get(False)})
fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.linear.default:lambda x,weight,bias=None:torch.ops.vds.linear(x,weight,bias)})

fx_infra.decomp.update_pre_convert_decomp({
    torch.ops.aten.linear.default:lambda x,weight,bias=None:torch.ops.vds.linear(x,weight,bias),
    torch.ops.aten.view.default:dynamic_view,
    torch.ops.aten._unsafe_view.default:dynamic_view,
    torch.ops.aten.native_layer_norm.default:dynamic_layer_norm,
    torch.ops.aten.arange.start_step:dynamic_arange,
    torch.ops.aten.expand.default:lambda x,shape,implicit=False:torch.ops.vds.broadcast_to(x,shape),
})

@torch.library.custom_op('vds::resize2d',mutates_args=(),schema='(Tensor x, SymInt[] size, bool align_corners, bool nearest) -> Tensor')
def dynamic_resize2d(x,size,align_corners,nearest):
    if nearest:return torch.nn.functional.interpolate(x,size=size,mode='nearest')
    return torch.nn.functional.interpolate(x,size=size,mode='bilinear',align_corners=align_corners)
@dynamic_resize2d.register_fake
def dynamic_resize2d_fake(x,size,align_corners,nearest):return x.new_empty((x.shape[0],x.shape[1],*size))
@registry.lower(torch.ops.vds.resize2d.default)
def lower_resize2d(lctx,x,size,align_corners,nearest):
    from litert_converter.mlir.dialects import stablehlo
    nhwc=stablehlo.transpose(x,[0,2,3,1])
    result=utils.node_meta_to_ir_types(lctx.node)[0]
    typ=ir.RankedTensorType.get([result.shape[i] for i in [0,2,3,1]],result.element_type)
    resized=_ir_operation('tfl.resize_nearest_neighbor' if nearest else 'tfl.resize_bilinear',results=[typ],operands=[nhwc,utils.convert_shape_to_ir_value(size)],attributes={'align_corners':ir.BoolAttr.get(align_corners),'half_pixel_centers':ir.BoolAttr.get(not align_corners and not nearest)})
    return stablehlo.transpose(resized,[0,3,1,2])
def bilinear_vec(x,output_size,align_corners,scale_factors=None):
    size=output_size if output_size is not None else [int(x.shape[-2+i]*scale_factors[i]) for i in range(2)]
    return torch.ops.vds.resize2d(x,size,align_corners,False)
def bilinear_default(x,output_size,align_corners,scales_h=None,scales_w=None):
    return torch.ops.vds.resize2d(x,output_size,align_corners,False)
fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.upsample_bilinear2d.vec:bilinear_vec,torch.ops.aten.upsample_bilinear2d.default:bilinear_default})

from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_mul_lowering
@registry.lower(torch.ops.tfl.mul.default)
def lower_bool_mul(lctx,lhs,rhs,fused_activation_function='NONE'):
    if isinstance(lhs,ir.Value) and isinstance(rhs,ir.Value) and str(lhs.type.element_type)=='i1' and str(rhs.type.element_type)=='i1':
        return _ir_operation('tfl.logical_and',results=utils.node_meta_to_ir_types(lctx.node),operands=[lhs,rhs])
    return _tfl_mul_lowering(lctx,lhs,rhs,fused_activation_function)

@torch.library.custom_op('vds::strided_slice',mutates_args=(),schema='(Tensor x, SymInt[] begin, SymInt[] end, int[] strides) -> Tensor')
def dynamic_strided_slice(x,begin,end,strides):
    return x[tuple(slice(b,e,s) for b,e,s in zip(begin,end,strides))].clone()
@dynamic_strided_slice.register_fake
def dynamic_strided_slice_fake(x,begin,end,strides):
    return x[tuple(slice(b,e,s) for b,e,s in zip(begin,end,strides))].clone()
@registry.lower(torch.ops.vds.strided_slice.default)
def lower_strided_slice(lctx,x,begin,end,strides):
    from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_strided_slice_lowering
    return _tfl_strided_slice_lowering(lctx,x,begin,end,strides)
def dynamic_slice(x,dim=0,start=None,end=None,step=1):
    n=x.shape[dim];b=0 if start is None else start;e=n if end is None else end
    if b<0:b+=n
    if e<0:e+=n
    b=max(0,min(b,n));e=max(b,min(e,n))
    begin=[0]*x.ndim;ends=list(x.shape);strides=[1]*x.ndim
    begin[dim]=b;ends[dim]=e;strides[dim]=step
    return torch.ops.vds.strided_slice(x,begin,ends,strides)
def dynamic_embedding(weight,indices,padding_idx=-1,scale_grad_by_freq=False,sparse=False):
    flat=torch.ops.vds.reshape(indices,[-1]).to(torch.int32)
    result=torch.ops.tfl.embedding_lookup(flat,weight)
    return torch.ops.vds.reshape(result,[*indices.shape,weight.shape[-1]])
fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.slice.Tensor:dynamic_slice,torch.ops.aten.embedding.default:dynamic_embedding})