Imports System
Imports System.Drawing
Imports Microsoft.Kinect
Imports System.Drawing.Imaging
Imports System.Runtime.InteropServices
Imports System.Windows
Imports System.IO
Imports Media = System.Windows.Media

Namespace ImageManipulationExtensionMethods
    Public Module ImageExtensions
        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(data() As Byte,
                                 width As Integer,
                                 height As Integer,
                                 format As PixelFormat) As Bitmap
            Dim _bitmap As New Bitmap(width, height, format)

            Dim _bitmapData As BitmapData = _bitmap.LockBits(
                New System.Drawing.Rectangle(0,
                                             0,
                                             _bitmap.Width,
                                             _bitmap.Height),
                                         ImageLockMode.WriteOnly,
                                         _bitmap.PixelFormat)
            Marshal.Copy(data, 0, _bitmapData.Scan0, data.Length)
            _bitmap.UnlockBits(_bitmapData)
            Return _bitmap
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(data() As Short,
                                 width As Integer,
                                 height As Integer,
                                 format As PixelFormat) As Bitmap
            Dim _bitmap As New Bitmap(width, height, format)

            Dim _bitmapData As BitmapData = _bitmap.LockBits(
                New System.Drawing.Rectangle(0,
                                            0,
                                            _bitmap.Width,
                                            _bitmap.Height),
                                        ImageLockMode.WriteOnly,
                                        _bitmap.PixelFormat)
            Marshal.Copy(data, 0, _bitmapData.Scan0, data.Length)
            _bitmap.UnlockBits(_bitmapData)
            Return _bitmap
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmapSource(data() As Byte,
                                       format As Media.PixelFormat,
                                       width As Integer,
                                       height As Integer) As Media.Imaging.BitmapSource
            Return Media.Imaging.BitmapSource.Create(width,
                                                     height,
                                                     96,
                                                     96,
                                                     format,
                                                     Nothing,
                                                     data,
                                                     width * format.BitsPerPixel / 8)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmapSource(data() As Short,
                                       format As Media.PixelFormat,
                                       width As Integer,
                                       height As Integer) As Media.Imaging.BitmapSource
            Return Media.Imaging.BitmapSource.Create(width,
                                                     height,
                                                     96,
                                                     96,
                                                     format,
                                                     Nothing,
                                                     data,
                                                     width * format.BitsPerPixel / 8)
        End Function

        ' bitmap methods

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(image As ColorImageFrame, format As PixelFormat) As Bitmap
            If (image Is Nothing OrElse image.PixelDataLength = 0) Then
                Return Nothing
            End If
            Dim data(image.PixelDataLength - 1) As Byte
            image.CopyPixelDataTo(data)
            Return data.ToBitmap(image.Width, image.Height, format)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(image As DepthImageFrame, format As PixelFormat) As Bitmap
            If (image Is Nothing OrElse image.PixelDataLength = 0) Then
                Return Nothing
            End If
            Dim data(image.PixelDataLength - 1) As Short
            image.CopyPixelDataTo(data)
            Return data.ToBitmap(image.Width, image.Height, format)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(image As ColorImageFrame) As Bitmap
            Return image.ToBitmap(PixelFormat.Format32bppRgb)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(image As DepthImageFrame) As Bitmap
            Return image.ToBitmap(PixelFormat.Format16bppRgb565)
        End Function

        ' bitmapsource methods

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmapSource(image As ColorImageFrame) As Media.Imaging.BitmapSource
            If (image Is Nothing OrElse image.PixelDataLength = 0) Then
                Return Nothing
            End If
            Dim data(image.PixelDataLength - 1) As Byte
            image.CopyPixelDataTo(data)
            Return data.ToBitmapSource(Media.PixelFormats.Bgr32, image.Width, image.Height)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmapSource(image As DepthImageFrame) As Media.Imaging.BitmapSource
            If (image Is Nothing OrElse image.PixelDataLength = 0) Then
                Return Nothing
            End If
            Dim data(image.PixelDataLength - 1) As Short
            image.CopyPixelDataTo(data)
            Return data.ToBitmapSource(Media.PixelFormats.Bgr555, image.Width, image.Height)
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToTransparentBitmapSource(data() As Byte,
                                                  width As Integer,
                                                  height As Integer) As Media.Imaging.BitmapSource
            Return data.ToBitmapSource(Media.PixelFormats.Bgra32, width, height)
        End Function


        ' conversion between bitmapsource and bitmap

        <DllImport("gdi32")> _
        Private Function DeleteObject(o As IntPtr) As Integer
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmapSource(_bitmap As Bitmap) As Media.Imaging.BitmapSource
            If (_bitmap Is Nothing) Then
                Return Nothing
            End If
            Dim ptr As IntPtr = _bitmap.GetHbitmap()
            Dim source As BitmapSource = System.Windows.Interop.Imaging.CreateBitmapSourceFromHBitmap(
                                ptr,
                                IntPtr.Zero,
                                Int32Rect.Empty,
                                Media.Imaging.BitmapSizeOptions.FromEmptyOptions())
            DeleteObject(ptr)
            Return source
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function ToBitmap(source As Media.Imaging.BitmapSource) As Bitmap
            Dim _bitmap As Bitmap
            Using outStream As MemoryStream = New MemoryStream()
                Dim enc As New Media.Imaging.PngBitmapEncoder()
                enc.Frames.Add(Media.Imaging.BitmapFrame.Create(source))
                enc.Save(outStream)
                _bitmap = New Bitmap(outStream)
            End Using
            Return _bitmap
        End Function

        <Runtime.CompilerServices.Extension()> _
        Public Function Flip(_bitmap As Bitmap) As Bitmap
            If (_bitmap Is Nothing) Then
                Return Nothing
            End If
            _bitmap.RotateFlip(RotateFlipType.RotateNoneFlipX)
            Return _bitmap
        End Function

    End Module
End Namespace